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

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

// Create an instance of the display
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Pin Definitions
#define TRIG_PIN1 5
#define ECHO_PIN1 18
#define TRIG_PIN2 19
#define ECHO_PIN2 23
#define RED_LED_PIN 25
#define YELLOW_LED_PIN 26
#define GREEN_LED_PIN 27
#define RELAY_PIN 14

// Variables
long duration1, duration2;
float distance1, distance2;

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

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_PIN, OUTPUT);
  pinMode(GREEN_LED_PIN, OUTPUT);
  pinMode(RELAY_PIN, OUTPUT);

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

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();
  displayInfo();
  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) {
    signalA = RED;
    digitalWrite(RED_LED_PIN, HIGH);
    digitalWrite(YELLOW_LED_PIN, LOW);
    digitalWrite(GREEN_LED_PIN, LOW);
    digitalWrite(RELAY_PIN, HIGH);
  } else if (distance1 > 300 && distance1 <= 700) {
    signalA = YELLOW;
    digitalWrite(RED_LED_PIN, LOW);
    digitalWrite(YELLOW_LED_PIN, HIGH);
    digitalWrite(GREEN_LED_PIN, LOW);
    digitalWrite(RELAY_PIN, LOW);
  } else {
    signalA = GREEN;
    digitalWrite(RED_LED_PIN, LOW);
    digitalWrite(YELLOW_LED_PIN, LOW);
    digitalWrite(GREEN_LED_PIN, HIGH);
    digitalWrite(RELAY_PIN, LOW);
  }
}

void displayInfo() {
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print("Train 501 Dist: ");
  display.print(distance1);
  display.println(" cm");
  
  display.setCursor(0, 10);
  display.print("Train 505 Dist: ");
  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");

  display.display();
}

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