#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <LiquidCrystal_I2C.h>
#include <ESP32Servo.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define SSD1306_I2C_ADDRESS 0x3C

#define LCD_ADDRESS 0x27
#define LCD_COLUMNS 20
#define LCD_ROWS 4

#define RELAY_PIN 32

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
LiquidCrystal_I2C lcd(LCD_ADDRESS, LCD_COLUMNS, LCD_ROWS);

#define TRIG_PIN1 5
#define ECHO_PIN1 18
#define TRIG_PIN2 19
#define ECHO_PIN2 23

#define RED_LED_PIN 25
#define YELLOW_LED_PIN1 26
#define YELLOW_LED_PIN2 33
#define GREEN_LED_PIN 27

#define SERVO_PIN 13

Servo gateServo;

float distance1, distance2;

enum SignalState { RED, YELLOW, GREEN };
SignalState signalA = GREEN, signalB = GREEN;

int axleCountStartTrain501 = 16;  // Total axles for Train 501
int axleCountEndTrain501 = 0;
int axleCountStartTrain505 = 20;  // Total axles for Train 505
int axleCountEndTrain505 = 0;

void setup() {
  Serial.begin(115200);

  pinMode(TRIG_PIN1, OUTPUT);
  pinMode(ECHO_PIN1, INPUT);
  pinMode(TRIG_PIN2, OUTPUT);
  pinMode(ECHO_PIN2, INPUT);

  pinMode(RED_LED_PIN, OUTPUT);
  pinMode(YELLOW_LED_PIN1, OUTPUT);
  pinMode(YELLOW_LED_PIN2, OUTPUT);
  pinMode(GREEN_LED_PIN, OUTPUT);

  pinMode(RELAY_PIN, OUTPUT);

  gateServo.attach(SERVO_PIN);

  if (!display.begin(SSD1306_I2C_ADDRESS, OLED_RESET)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);

  lcd.begin(LCD_COLUMNS, LCD_ROWS);
  lcd.backlight();

  randomSeed(analogRead(0));
}

void loop() {
  // Simulate random distances for trains
  distance1 = random(0, 1501);
  distance2 = random(0, 1501);

  // Update axle counts based on distance
  axleCountEndTrain501 = updateEndCount(axleCountStartTrain501, distance1);
  axleCountEndTrain505 = updateEndCount(axleCountStartTrain505, distance2);

  Serial.print("Train 501 Distance: ");
  Serial.print(distance1);
  Serial.println(" cm");

  Serial.print("Train 505 Distance: ");
  Serial.print(distance2);
  Serial.println(" cm");

  updateSignals();
  controlRelay();

  displayInfo();
  lcdInfo();

  delay(1000);
}

int updateEndCount(int startCount, float distance) {
  // Define distance thresholds
  const float TRAIN_LENGTH = 1500;  // Total length of the track section in cm where the train is considered fully passed

  // Assume that as the train moves beyond the TRAIN_LENGTH, it has fully passed
  if (distance > TRAIN_LENGTH) {
    return 0;  // All axles have passed
  } else {
    // Calculate the proportion of the train length that has passed based on distance
    int axlesPassing = map(distance, 0, TRAIN_LENGTH, 0, startCount);
    return max(startCount - axlesPassing, 0); // Ensure end count is not negative
  }
}

void updateSignals() {
  if (distance1 <= 300 || distance2 <= 300) {
    signalA = RED;
    signalB = RED;
    setSignalState(RED);
    gateServo.write(0);
  } else if ((distance1 > 300 && distance1 <= 700) || (distance2 > 300 && distance2 <= 700)) {
    signalA = YELLOW;
    signalB = RED;
    setSignalState(YELLOW);
    gateServo.write(45);
  } else if ((distance1 > 700 && distance1 <= 1500) || (distance2 > 700 && distance2 <= 1500)) {
    signalA = GREEN;
    signalB = YELLOW;
    setSignalState(GREEN);
    gateServo.write(90);
  } else {
    signalA = GREEN;
    signalB = GREEN;
    setSignalState(GREEN);
    gateServo.write(90);
  }
}

void setSignalState(SignalState state) {
  digitalWrite(RED_LED_PIN, state == RED);
  digitalWrite(YELLOW_LED_PIN1, state == YELLOW);
  digitalWrite(YELLOW_LED_PIN2, state == YELLOW);
  digitalWrite(GREEN_LED_PIN, state == GREEN);
}

void controlRelay() {
  if (distance1 <= 300) {
    digitalWrite(RELAY_PIN, HIGH);
    Serial.println("Train 501 Stopping: Relay Activated");
    updateDisplayAndLCD("Train 501 Stopping", "Train 505 Moving");
  } else if (distance2 <= 300) {
    digitalWrite(RELAY_PIN, HIGH);
    Serial.println("Train 505 Stopping: Relay Activated");
    updateDisplayAndLCD("Train 505 Stopping", "Train 501 Moving");
  } else {
    digitalWrite(RELAY_PIN, LOW);
    Serial.println("Both Trains Moving: Relay Deactivated");
    updateDisplayAndLCD("Train 501 Moving", "Train 505 Moving");
  }
}

void updateDisplayAndLCD(const char* train1Status, const char* train2Status) {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print(train1Status);
  display.setCursor(0, 10);
  display.print(train2Status);
  display.display();

  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print(train1Status);
  lcd.setCursor(0, 1);
  lcd.print(train2Status);
}

void displayInfo() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("Train 501: ");
  display.println(axleCountStartTrain501 == axleCountEndTrain501 ? "Passed" : "Not Passed");
  display.setCursor(0, 10);
  display.print("Train 505: ");
  display.println(axleCountStartTrain505 == axleCountEndTrain505 ? "Passed" : "Not Passed");
  display.setCursor(0, 20);
  display.print("Signal A: ");
  display.println(signalToString(signalA));
  display.setCursor(0, 30);
  display.print("Signal B: ");
  display.println(signalToString(signalB));
  display.setCursor(0, 40);
  display.print("Time: ");
  display.print(millis() / 1000);
  display.println(" s");
  display.display();
}

void lcdInfo() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("501 Start: ");
  lcd.print(axleCountStartTrain501);
  lcd.print(" End: ");
  lcd.print(axleCountEndTrain501);
  lcd.setCursor(0, 1);
  lcd.print(axleCountStartTrain501 == axleCountEndTrain501 ? "501 Passed" : "501 Not Passed");

  lcd.setCursor(0, 2);
  lcd.print("505 Start: ");
  lcd.print(axleCountStartTrain505);
  lcd.print(" End: ");
  lcd.print(axleCountEndTrain505);
  lcd.setCursor(0, 3);
  lcd.print(axleCountStartTrain505 == axleCountEndTrain505 ? "505 Passed" : "505 Not Passed");
}

String signalToString(SignalState signal) {
  switch (signal) {
    case RED: return "RED";
    case YELLOW: return "YELLOW";
    case GREEN: return "GREEN";
    default: return "UNKNOWN";
  }
}