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

// OLED Configuration
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Traffic Light LEDs (Vehicles)
#define DRIVER_RED 4
#define DRIVER_YELLOW 16
#define DRIVER_GREEN 17

// Pedestrian Light LEDs
#define PED_RED 14
#define PED_GREEN 12

// Ultrasonic Sensor 1 (Vehicle Detection)
#define TRIG_1 27
#define ECHO_1 26

// Ultrasonic Sensor 2 (Speed Measurement)
#define TRIG_2 18
#define ECHO_2 19

// Variables
long vehicleDistance1 = 0;
long vehicleDistance2 = 0;
float speed = 0;
unsigned long prevTime = 0;

// Function to get distance from Ultrasonic Sensor
long getDistance(int trigPin, int echoPin) {
    digitalWrite(trigPin, LOW);
    delayMicroseconds(2);
    digitalWrite(trigPin, HIGH);
    delayMicroseconds(10);
    digitalWrite(trigPin, LOW);
    return pulseIn(echoPin, HIGH) * 0.034 / 2;  // Convert to cm
}

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

    // Initialize OLED Display
    if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
        Serial.println("SSD1306 allocation failed");
        while (1);
    }
    
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(WHITE);

    // Initialize LED Pins
    pinMode(DRIVER_RED, OUTPUT);
    pinMode(DRIVER_YELLOW, OUTPUT);
    pinMode(DRIVER_GREEN, OUTPUT);
    pinMode(PED_RED, OUTPUT);
    pinMode(PED_GREEN, OUTPUT);

    // Initialize Ultrasonic Sensors
    pinMode(TRIG_1, OUTPUT);
    pinMode(ECHO_1, INPUT);
    pinMode(TRIG_2, OUTPUT);
    pinMode(ECHO_2, INPUT);

    // Start Traffic Light with Green
    digitalWrite(DRIVER_GREEN, HIGH);  // Vehicles move
    digitalWrite(DRIVER_YELLOW, LOW);
    digitalWrite(DRIVER_RED, LOW);
    digitalWrite(PED_RED, HIGH);  // Pedestrians stop
    digitalWrite(PED_GREEN, LOW);
}

void loop() {
    // Read vehicle distances
    vehicleDistance1 = getDistance(TRIG_1, ECHO_1);
    vehicleDistance2 = getDistance(TRIG_2, ECHO_2);

    Serial.print("Vehicle Distance 1: ");
    Serial.print(vehicleDistance1);
    Serial.print(" cm | Distance 2: ");
    Serial.print(vehicleDistance2);
    Serial.print(" cm | Speed: ");
    Serial.println(speed);

    // Calculate speed if a vehicle is detected at both sensors
    if (vehicleDistance1 > 0 && vehicleDistance1 < 100 && vehicleDistance2 > 0 && vehicleDistance2 < 100) {
        unsigned long currentTime = millis();
        if (prevTime != 0) {
            float timeDiff = (currentTime - prevTime) / 1000.0;  // Convert ms to seconds
            float distanceBetweenSensors = 10.0;  // Assume 10 cm between sensors
            speed = distanceBetweenSensors / timeDiff;  // Speed in cm/s
        }
        prevTime = currentTime;
    } else {
        speed = 0;  // No vehicle detected
    }

    // Run automatic traffic light cycle
    automaticTrafficCycle();
    
    updateOLED();

    delay(500);
}

void automaticTrafficCycle() {
    // Green Light for Vehicles (Pedestrians Wait)
    Serial.println("Green Light: Vehicles move");
    digitalWrite(DRIVER_GREEN, HIGH);
    digitalWrite(DRIVER_YELLOW, LOW);
    digitalWrite(DRIVER_RED, LOW);
    digitalWrite(PED_RED, HIGH);
    digitalWrite(PED_GREEN, LOW);
    delay(5000);

    // Yellow Light (Warning)
    Serial.println("Yellow Light: Prepare to stop");
    digitalWrite(DRIVER_GREEN, LOW);
    digitalWrite(DRIVER_YELLOW, HIGH);
    delay(3000);

    // Red Light for Vehicles, Green for Pedestrians
    Serial.println("Red Light: Vehicles stop, Pedestrians cross");
    digitalWrite(DRIVER_YELLOW, LOW);
    digitalWrite(DRIVER_RED, HIGH);
    digitalWrite(PED_RED, LOW);
    digitalWrite(PED_GREEN, HIGH);
    delay(5000);

    // Back to Green for Vehicles
    Serial.println("Back to Green Light for Vehicles");
    digitalWrite(DRIVER_RED, LOW);
    digitalWrite(DRIVER_GREEN, HIGH);
    digitalWrite(PED_GREEN, LOW);
    digitalWrite(PED_RED, HIGH);
}

void updateOLED() {
    display.clearDisplay();
    display.setCursor(0, 0);

    // Display vehicle detection status
    if (vehicleDistance1 > 0 && vehicleDistance1 < 100) {
        display.println("Vehicle Detected!");
    } else {
        display.println("No Vehicle");
    }

    // Display Speed
    display.print("Speed: ");
    display.print(speed);
    display.println(" cm/s");

    display.display();
}