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

// OLED display parameters
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define SSD1306_I2C_ADDRESS 0x3C

// LCD 20x4 I2C parameters
#define LCD_ADDRESS 0x27
#define LCD_COLUMNS 20
#define LCD_ROWS 4

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 BUZZER_PIN 32
#define SERVO_PIN 13

long duration1, duration2;
float distance1, distance2;

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

int wheelCount = 0;
int knownWheels = 8;
bool trainDetected = false;

Servo gateServo;

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(BUZZER_PIN, OUTPUT);

  gateServo.attach(SERVO_PIN);

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

  // Initialize LCD with number of columns, rows, and default character size
  lcd.begin(LCD_COLUMNS, LCD_ROWS);
  lcd.backlight();
}

void loop() {
  distance1 = getDistance(TRIG_PIN1, ECHO_PIN1);
  distance2 = getDistance(TRIG_PIN2, ECHO_PIN2);

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

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

  updateSignals();
  checkTrainPassing();  

  displayInfo();
  lcdInfo();
  delay(1000);
}

float getDistance(int trigPin, int echoPin) {
  digitalWrite(trigPin, LOW);
  delayMicroseconds(2);
  digitalWrite(trigPin, HIGH);
  delayMicroseconds(10);
  digitalWrite(trigPin, LOW);

  long duration = pulseIn(echoPin, HIGH);
  return (duration * 0.0343) / 2;
}

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);  // Close gate
    tone(BUZZER_PIN, 1000); // Sound buzzer
  } 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);  // Gate in transit position
    noTone(BUZZER_PIN);
  } 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);  // Gate open
    noTone(BUZZER_PIN);
  } 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);  // Gate open
    noTone(BUZZER_PIN);
  }
}

void checkTrainPassing() {
  if (distance1 < 10) {
    if (!trainDetected) {
      trainDetected = true;
      wheelCount++;
      Serial.print("Wheel Count: ");
      Serial.println(wheelCount);
    }
  } else {
    trainDetected = false;
  }

  if (wheelCount == knownWheels) {
    Serial.println("Train has passed completely.");
    wheelCount = 0;  // Reset for the next train
  }
}

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.println("Train 501 in Block 1");
  } else if (distance1 > 300 && distance1 <= 700) {
    display.println("Train 501 in Block 2");
  } else if (distance1 > 700 && distance1 <= 1500) {
    display.println("Train 501 in Block 3");
  } else {
    display.println("Train 501 is arriving at Station B");
  }

  display.setCursor(0, 60);
  if (distance2 <= 300) {
    display.println("Train 505 in Block 1");
  } else if (distance2 > 300 && distance2 <= 700) {
    display.println("Train 505 in Block 2");
  } else if (distance2 > 700 && distance2 <= 1500) {
    display.println("Train 505 in Block 3");
  } else {
    display.println("Train 505 is arriving at Station B");
  }
}

void lcdInfo() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("Train 501 Dist: ");
  lcd.print(distance1);
  lcd.print(" cm");

  lcd.setCursor(0, 1);
  lcd.print("Train 505 Dist: ");
  lcd.print(distance2);
  lcd.print(" cm");

  lcd.setCursor(0, 2);
  lcd.print("Signal A: ");
  lcd.print(signalToString(signalA));

  lcd.setCursor(0, 3);
  lcd.print("Signal B: ");
  lcd.print(signalToString(signalB));
}

const char* signalToString(SignalState state) {
  switch (state) {
    case RED: return "Red";
    case YELLOW: return "Yellow";
    case GREEN: return "Green";
    default: return "Unknown";
  }
}