#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <ESP32Servo.h>
#include <LiquidCrystal_I2C.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 wheelCountTrain501 = 0;
int wheelCountTrain505 = 0;
int expectedBogiesTrain501 = 12; // 12 bogies * 2 wheels = 24 wheels
int expectedBogiesTrain505 = 12; // 12 bogies * 2 wheels = 24 wheels

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() {
  // Generate random distances for simulation
  distance1 = random(0, 1501);
  distance2 = random(0, 1501);

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

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

  // Update the signal states based on the distances
  updateSignals();

  // Update the train passing statuses based on distances
  checkTrainPassingStatus(distance1, expectedBogiesTrain501, "Train 501", wheelCountTrain501);
  checkTrainPassingStatus(distance2, expectedBogiesTrain505, "Train 505", wheelCountTrain505);

  // Control the relay based on the distances
  controlRelay();

  // Update OLED and LCD displays with current information
  displayInfo();
  lcdInfo();
  delay(1000);
}

void updateSignals() {
  if (distance1 <= 300 || distance2 <= 300) {
    signalA = RED;
    signalB = RED;
    digitalWrite(RED_LED_PIN, HIGH);
    digitalWrite(YELLOW_LED_PIN1, LOW);
    digitalWrite(YELLOW_LED_PIN2, LOW);
    digitalWrite(GREEN_LED_PIN, LOW);
    gateServo.write(0);
  } else if ((distance1 > 300 && distance1 <= 700) || (distance2 > 300 && distance2 <= 700)) {
    signalA = YELLOW;
    signalB = RED;
    digitalWrite(RED_LED_PIN, LOW);
    digitalWrite(YELLOW_LED_PIN1, HIGH);
    digitalWrite(YELLOW_LED_PIN2, LOW);
    digitalWrite(GREEN_LED_PIN, LOW);
    gateServo.write(45);
  } else if ((distance1 > 700 && distance1 <= 1500) || (distance2 > 700 && distance2 <= 1500)) {
    signalA = GREEN;
    signalB = YELLOW;
    digitalWrite(RED_LED_PIN, LOW);
    digitalWrite(YELLOW_LED_PIN1, LOW);
    digitalWrite(YELLOW_LED_PIN2, HIGH);
    digitalWrite(GREEN_LED_PIN, HIGH);
    gateServo.write(90);
  } else {
    signalA = GREEN;
    signalB = GREEN;
    digitalWrite(RED_LED_PIN, LOW);
    digitalWrite(YELLOW_LED_PIN1, LOW);
    digitalWrite(YELLOW_LED_PIN2, LOW);
    digitalWrite(GREEN_LED_PIN, HIGH);
    gateServo.write(90);
  }
}

void checkTrainPassingStatus(float distance, int expectedBogies, const char* trainName, int &wheelCount) {
  const int WHEELS_PER_BOGIE = 2; // Visible wheels per bogie to the sensor
  const float WHEEL_DETECTION_THRESHOLD = 10.0;

  // Check for wheel detection based on distance
  if (distance <= WHEEL_DETECTION_THRESHOLD) {
    wheelCount++;
  }

  // Calculate number of bogies detected
  int detectedBogies = wheelCount / WHEELS_PER_BOGIE;

  // Display train passing status on Serial and LCD
  if (detectedBogies == expectedBogies) {
    Serial.print(trainName);
    Serial.println(" Passed");
    lcd.setCursor(0, 0);
    lcd.print(trainName);
    lcd.print(" Passed");
  } else if (detectedBogies < expectedBogies) {
    Serial.print(trainName);
    Serial.println(" Not Passed");
    lcd.setCursor(0, 0);
    lcd.print(trainName);
    lcd.print(" Not Passed");
  } else {
    Serial.print(trainName);
    Serial.println(" Error");
    lcd.setCursor(0, 0);
    lcd.print(trainName);
    lcd.print(" Error");
  }

  // Reset wheel count for next train if passed
  if (detectedBogies >= expectedBogies) {
    wheelCount = 0;
  }
}

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 Distance: ");
  display.print(distance1);
  display.println(" cm");
  display.setCursor(0, 10);
  display.print("Train 505 Distance: ");
  display.print(distance2);
  display.println(" cm");
  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");

  displayTrainStatus();
  display.display();
}

void displayTrainStatus() {
  display.setCursor(0, 50);
  if (distance1 <= 300) {
    display.print("Train 501 at Station B - Block 1");
  } else if (distance1 > 300 && distance1 <= 700) {
    display.print("Train 501 at Station B - Block 2");
  } else if (distance1 > 700 && distance1 <= 1500) {
    display.print("Train 501 approaching Station B");
  } else {
    display.print("Train 501 leaving Station B");
  }

  display.setCursor(0, 60);
  if (distance2 <= 300) {
    display.print("Train 505 at Station B - Block 1");
  } else if (distance2 > 300 && distance2 <= 700) {
    display.print("Train 505 at Station B - Block 2");
  } else if (distance2 > 700 && distance2 <= 1500) {
    display.print("Train 505 approaching Station B");
  } else {
    display.print("Train 505 leaving Station B");
  }
}

void lcdInfo() {
  lcd.clear();
  
  // Display Train 501 information
  lcd.setCursor(0, 0);
  lcd.print("Train 501 Dist: ");
  lcd.print(distance1);
  lcd.print(" cm");
  
  lcd.setCursor(0, 1);
  lcd.print("Block: ");
  if (distance1 <= 300) {
    lcd.print("Station B - Block 1");
  } else if (distance1 > 300 && distance1 <= 700) {
    lcd.print("Station B - Block 2");
  } else if (distance1 > 700 && distance1 <= 1500) {
    lcd.print("Approaching");
  } else {
    lcd.print("Leaving");
  }
  
  // Display Train 505 information
  lcd.setCursor(0, 2);
  lcd.print("Train 505 Dist: ");
  lcd.print(distance2);
  lcd.print(" cm");
  
  lcd.setCursor(0, 3);
  lcd.print("Block: ");
  if (distance2 <= 300) {
    lcd.print("Station B - Block 1");
  } else if (distance2 > 300 && distance2 <= 700) {
    lcd.print("Station B - Block 2");
  } else if (distance2 > 700 && distance2 <= 1500) {
    lcd.print("Approaching");
  } else {
    lcd.print("Leaving");
  }
}

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