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

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

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Pin definitions
#define GREEN_LED 25
#define YELLOW_LED 26
#define RED_LED 27
#define TRIG_PIN 5
#define ECHO_PIN 18
#define BUZZER_PIN 4

long duration;
float distance;

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

  // LED pins
  pinMode(GREEN_LED, OUTPUT);
  pinMode(YELLOW_LED, OUTPUT);
  pinMode(RED_LED, OUTPUT);

  // Ultrasonic sensor
  pinMode(TRIG_PIN, OUTPUT);
  pinMode(ECHO_PIN, INPUT);

  // Buzzer
  pinMode(BUZZER_PIN, OUTPUT);

  // OLED initialization
  if (!display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDRESS)) {
    Serial.println("SSD1306 allocation failed");
    for (;;);
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(10, 10);
  display.println("Traffic System Init...");
  display.display();

  delay(1000);
}

float getDistance() {
  digitalWrite(TRIG_PIN, LOW);
  delayMicroseconds(2);

  digitalWrite(TRIG_PIN, HIGH);
  delayMicroseconds(10);

  digitalWrite(TRIG_PIN, LOW);

  duration = pulseIn(ECHO_PIN, HIGH);

  float dist = duration * 0.034 / 2;

  return dist;
}

void showNormalDisplay(String signal) {
  display.clearDisplay();

  display.setTextSize(2);
  display.setTextColor(WHITE);
  display.setCursor(10, 20);

  display.print(signal);

  display.display();
}

void redLightMonitor() {

  unsigned long startTime = millis();

  while (millis() - startTime < 10000) {

    digitalWrite(RED_LED, HIGH);

    distance = getDistance();

    Serial.print("Distance: ");
    Serial.print(distance);
    Serial.println(" cm");

    // Vehicle detected within 10 cm
    if (distance > 0 && distance < 10) {

      // Start buzzer
      tone(BUZZER_PIN, 2000);

      // OLED warning
      display.clearDisplay();
      display.setTextSize(1);
      display.setTextColor(WHITE);

      display.setCursor(5, 20);
      display.println("RED LIGHT");

      display.setCursor(5, 35);
      display.println("CROSSED!");

      display.display();

    } 
    else {

      // Stop buzzer
      noTone(BUZZER_PIN);

      // Normal RED display
      showNormalDisplay("RED");
    }

    delay(100);
  }

  // End of RED phase
  digitalWrite(RED_LED, LOW);
  noTone(BUZZER_PIN);
}

void loop() {

  // =========================
  // GREEN LIGHT
  // =========================

  digitalWrite(GREEN_LED, HIGH);

  showNormalDisplay("GREEN");

  delay(10000);

  digitalWrite(GREEN_LED, LOW);


  // =========================
  // YELLOW LIGHT
  // =========================

  digitalWrite(YELLOW_LED, HIGH);

  showNormalDisplay("YELLOW");

  delay(10000);

  digitalWrite(YELLOW_LED, LOW);


  // =========================
  // RED LIGHT
  // =========================

  redLightMonitor();
}