#include <LiquidCrystal.h>

// Pin constants
const int SWITCH_PIN = 21; 
const int BTN_1_PIN = 27;
const int BTN_2_PIN = 26;
const int BTN_3_PIN = 22;

const int LED_3_R_PIN = 6;   // RGB1 (Red)   - Connected to GP6
const int LED_3_G_PIN = 13;  // RGB1 (Green) - Connected to GP13
const int LED_3_B_PIN = 14;  // RGB1 (Blue)  - Connected to GP14

const int LED_2_R_PIN = 3;   // RGB2 (Red)   - Connected to GP3
const int LED_2_G_PIN = 4;   // RGB2 (Green) - Connected to GP4
const int LED_2_B_PIN = 5;   // RGB2 (Blue)  - Connected to GP5

const int LED_1_R_PIN = 0;   // RGB3 (Red)   - Connected to GP0
const int LED_1_G_PIN = 1;   // RGB3 (Green) - Connected to GP1
const int LED_1_B_PIN = 2;   // RGB3 (Blue)  - Connected to GP2

// Initialize lcd
LiquidCrystal lcd(12, 11, 10, 9, 8, 7);

// Enums
enum Mode {
  RANDOM,
  MANUAL
};

enum JunctionState { // Enum for junction state
    GREEN_1, // Lane 1 has green light
    GREEN_2, // Lane 2 has green light
    GREEN_3, // Lane 3 has green light
    YELLOW_1_TO_2, // Transitioning from 1 to 2
    YELLOW_2_TO_3, // Transitioning from 2 to 3
    YELLOW_3_TO_1  // Transitioning from 3 to 1
};

enum Color { // Enum for each traffic light's color
    RED,
    YELLOW,
    GREEN
};

// Global variables
Mode mode;
Mode prev_mode = MANUAL;
bool display_changed = true;
bool btn_pressed[4] = {false, false, false, false};
bool btn_prev[4] = {false, false, false, false};
enum MarkovState { STATE_A, STATE_B };
MarkovState lane_state[4] = {STATE_A, STATE_A, STATE_A, STATE_A}; // Index 0 unused, lanes 1-3
unsigned long last_car_time[4] = {0, 0, 0, 0}; // Last time a car was generated for each lane
unsigned long waiting_time[4] = {0, 0, 0, 0}; // Waiting time for next car in STATE_B
int next_car_id = 100; // Starting car ID
unsigned long last_debounce_time[4] = {0, 0, 0, 0}; // Time when button state last changed
const unsigned long debounce_delay = 50; // Debounce time in milliseconds (50ms is usually sufficient)
unsigned long last_departure_time = 0;
int green_duration = 0;  // Tracks how long the current light has been green
unsigned long last_optimization_time = 0;    // Tracks when we last ran optimization
const unsigned long OPTIMIZATION_INTERVAL = 30000;  // 30 seconds in milliseconds
int total_thrown_cars = 0; // Total thrown cars in current light pattern
bool has_completed_round = false; // Flag to track if a complete round has been completed
bool pattern_update_pending = false; // Flag to indicate a pattern update is waiting

// Structs and functions
struct Car;
struct Car {
    int car_ID;     // Unique ID for each car
    unsigned long arrival_time; // Time of arrival for wait time calculations
    Car* next;      // Pointer to next car in lane
    
    // Constructor
    Car(int id) : car_ID(id), arrival_time(millis()), next(nullptr) {
        display_changed = true; // When a car is generated - screen refresh is needed
    }
};

struct Lane {
    Car* head; // Pointer to first car in lane
    Car* tail; // Pointer to last car in lane 
    int num_of_cars; // Number of cars in the lane
    int thrown_cars; // Count of thrown cars for this lane in the current pattern
    
    // Constructor
    Lane() : head(nullptr), tail(nullptr), num_of_cars(0), thrown_cars(0) {}
    
    // Function to add a car to the end of the linked list (FIFO)
    void add_car(Car* car) {
        if (car == nullptr) return;
        
        car->next = nullptr; // Ensure the new car doesn't point to anything
        
        if (head == nullptr) {
            // If the lane is empty, the new car becomes both head and tail
            head = car;
            tail = car;
        } else {
            // Otherwise, add the car to the end and update the tail
            tail->next = car;
            tail = car;
        }
        
        num_of_cars++;
    }
    
    Car* sub_car() {
        if (head == nullptr) {
            // Lane is empty
            return nullptr;
        }
        
        Car* temp = head;  // Store the head to return it
        head = head->next; // Update the head to the next car
        
        if (head == nullptr) {
            // If the lane is now empty, update the tail as well
            tail = nullptr;
        }
        
        temp->next = nullptr; // Isolate the removed car
        num_of_cars--;
        
        return temp;
    }
};

struct TrafficLights {
    JunctionState current_state; // Current state of the junction
    Color TL_color[4]; // Colors for each traffic light (1-3)
    int green_duration[4]; // Green light duration for each TL (1-3)
    unsigned long last_state_change; // Time of the last junction state change
    bool yellow_transition; // Flag to track if we're in a yellow transition
    unsigned long yellow_start_time; // When the yellow light started
    JunctionState next_state; // The state we're transitioning to
    bool round_complete; // Whether a complete round has been completed
    bool pattern_changed; // Whether the pattern has changed and needs to be updated
    
    // Constructor
    TrafficLights() : 
        current_state(GREEN_1), 
        last_state_change(0), 
        yellow_transition(false),
        yellow_start_time(0),
        next_state(GREEN_1),
        round_complete(false),
        pattern_changed(false) 
    {
        TL_color[1] = GREEN;
        TL_color[2] = RED;
        TL_color[3] = RED;
        green_duration[1] = 3;
        green_duration[2] = 3;
        green_duration[3] = 3;
    }

    // Function to start yellow transition
    void start_yellow_transition(JunctionState new_state) {
        // Determine which traffic light is currently green
        int current_tl = 0;
        switch (current_state) {
            case GREEN_1: current_tl = 1; break;
            case GREEN_2: current_tl = 2; break;
            case GREEN_3: current_tl = 3; break;
            default: return; // Invalid state or already in yellow transition
        }
        
        // Set the light to yellow
        TL_color[current_tl] = YELLOW;
        
        // Log the traffic light change event
        Serial.print(millis());
        Serial.print(", system, Traffic Light ");
        Serial.print(current_tl);
        Serial.println(" changing from Green to Yellow");
        
        // Start yellow transition
        yellow_transition = true;
        yellow_start_time = millis();
        next_state = new_state;
        
        // Update display
        display_changed = true;
    }
    
    // Function to complete the transition after yellow period
    void complete_transition() {
        // Yellow period is over, move to the next state
        yellow_transition = false;
        
        // Update the traffic light colors based on next state
        for (int i = 1; i <= 3; i++) {
            TL_color[i] = RED; // Default all to red
        }
        
        // Set the new green light
        int new_green_tl = 0;
        switch (next_state) {
            case GREEN_1: new_green_tl = 1; break;
            case GREEN_2: new_green_tl = 2; break;
            case GREEN_3: new_green_tl = 3; break;
            default: return; // Invalid state
        }
        
        TL_color[new_green_tl] = GREEN;
        
        // Log the final color change
        Serial.print(millis());
        Serial.print(", system, Traffic Light ");
        Serial.print(new_green_tl);
        Serial.println(" changing from Red to Green");
        
        // Update the current state and record the time of change
        current_state = next_state;
        last_state_change = millis();
        
        // Check if we've completed a round (all TLs have been green at least once)
        if (next_state == GREEN_1) {
            round_complete = true;
            
            // If there's a pattern update pending, apply it now
            if (pattern_changed) {
                Serial.print(millis());
                Serial.print(", system, Light Pattern Applied: ");
                Serial.print(green_duration[1]);
                Serial.print("-");
                Serial.print(green_duration[2]);
                Serial.print("-");
                Serial.println(green_duration[3]);
                pattern_changed = false;
            }
        }
        
        // Update display
        display_changed = true;
    }
    
    // Function to check and process yellow transition
    void process_yellow_transition() {
        if (yellow_transition && (millis() - yellow_start_time >= 500)) {
            complete_transition();
        }
    }
    
    // Function to update the light pattern (only takes effect after a complete round)
    void update_pattern(int duration1, int duration2, int duration3) {
        green_duration[1] = duration1;
        green_duration[2] = duration2;
        green_duration[3] = duration3;
        pattern_changed = true;
        
        // Update display
        display_changed = true;
    }
};

struct Statistics {
    const static int HISTORY_SIZE = 60;
    int lane_length_history[4][HISTORY_SIZE]; // Index 0 unused, lanes 1-3
    int history_index;
    int thrown_cars[4]; // Index 0 unused, lanes 1-3
    unsigned long wait_time_sum[4]; // Total wait time for all cars in each lane
    int car_count[4]; // Number of cars processed for average calculation
    
    Statistics() {
        for (int lane = 0; lane < 4; lane++) {
            for (int i = 0; i < HISTORY_SIZE; i++) {
                lane_length_history[lane][i] = 0;
            }
            thrown_cars[lane] = 0;
            wait_time_sum[lane] = 0;
            car_count[lane] = 0;
        }
        history_index = 0;
    }
    
    void add_lane_length(int lane, int length) {
        if (lane >= 1 && lane <= 3) {
            lane_length_history[lane][history_index] = length;
        }
        
        // If we've updated lane 3, move to the next index
        if (lane == 3) {
            history_index = (history_index + 1) % HISTORY_SIZE;
        }
    }
    
    float calc_avg_lane_length(int lane) {
        if (lane < 1 || lane > 3) return 0.0;
        
        int sum = 0;
        for (int i = 0; i < HISTORY_SIZE; i++) {
            sum += lane_length_history[lane][i];
        }
        return (float)sum / HISTORY_SIZE;
    }
    
    void increment_thrown(int lane) {
        if (lane >= 1 && lane <= 3) {
            thrown_cars[lane]++;
        }
    }
    
    void add_wait_time(int lane, unsigned long wait_time) {
        if (lane >= 1 && lane <= 3) {
            wait_time_sum[lane] += wait_time;
            car_count[lane]++;
        }
    }
    
    float calc_avg_wait_time(int lane) {
        if (lane < 1 || lane > 3 || car_count[lane] == 0) return 0.0;
        return (float)wait_time_sum[lane] / car_count[lane];
    }
    
    int get_thrown(int lane) {
        if (lane >= 1 && lane <= 3) {
            return thrown_cars[lane];
        }
        return 0;
    }
    
    int get_total_thrown() {
        return thrown_cars[1] + thrown_cars[2] + thrown_cars[3];
    }
    
    void reset_thrown_cars() {
        for (int lane = 1; lane <= 3; lane++) {
            thrown_cars[lane] = 0;
        }
    }
};

// Initialize structures
Lane lane1, lane2, lane3;
TrafficLights junction;
Statistics stats;

// Setup functions
void setup_lcd() {
    lcd.begin(20, 4);
    lcd.clear();
}

void setup_serial() {
    Serial.begin(57600);
    Serial.println("Serial initialized");
}

void pin_modes_setup() {
    pinMode(SWITCH_PIN, INPUT_PULLUP);
    pinMode(BTN_1_PIN, INPUT_PULLDOWN);
    pinMode(BTN_2_PIN, INPUT_PULLDOWN);
    pinMode(BTN_3_PIN, INPUT_PULLDOWN);
    
    pinMode(LED_1_R_PIN, OUTPUT);
    pinMode(LED_2_R_PIN, OUTPUT);
    pinMode(LED_3_R_PIN, OUTPUT);
    
    pinMode(LED_1_G_PIN, OUTPUT);
    pinMode(LED_2_G_PIN, OUTPUT);
    pinMode(LED_3_G_PIN, OUTPUT);
    
    pinMode(LED_1_B_PIN, OUTPUT);
    pinMode(LED_2_B_PIN, OUTPUT);
    pinMode(LED_3_B_PIN, OUTPUT);
}

// Loop functions
void update_monitor() {
    lcd.clear();
    
    // Display mode
    lcd.setCursor(0, 0);
    lcd.print("MODE: ");
    (mode == MANUAL) ? lcd.print("MANUAL") : lcd.print("RANDOM");
    
    // Display TL colors
    lcd.setCursor(0, 1);
    for (int i = 1; i <= 3; i++) {
        lcd.print("TL");
        lcd.print(i);
        lcd.print(":");
        switch (junction.TL_color[i]) {
            case RED:
                lcd.print("R ");
                break;
            case YELLOW:
                lcd.print("Y ");
                break;
            case GREEN:
                lcd.print("G ");
                break;
        }
    }
    
    // Print lane occupancy
    lcd.setCursor(0, 2);
    for (int i = 1; i <= 3; i++) {
        lcd.print("L");
        lcd.print(i);
        lcd.print(":");

        switch (i) {
            case 1:
                lcd.print(lane1.num_of_cars);
                break;
            case 2:
                lcd.print(lane2.num_of_cars);
                break;
            case 3:
                lcd.print(lane3.num_of_cars);
                break;
        }
        lcd.print(" ");
    }
    
    // Display light pattern and thrown cars
    lcd.setCursor(0, 3);
    lcd.print("Pat:");
    lcd.print(junction.green_duration[1]);
    lcd.print("-");
    lcd.print(junction.green_duration[2]);
    lcd.print("-");
    lcd.print(junction.green_duration[3]);
    lcd.print(" Thr:");
    lcd.print(total_thrown_cars);
}

void update_inputs() {
    Mode current_mode = digitalRead(SWITCH_PIN) ? MANUAL : RANDOM;
    
    // Check if mode has changed
    if (current_mode != prev_mode) {
        display_changed = true;
        Serial.print("Mode changed to: ");
        Serial.println(current_mode == MANUAL ? "MANUAL" : "RANDOM");
    }
    
    // Update the mode
    mode = current_mode;
    prev_mode = current_mode;
}

void update_leds() {
    for (int i = 1; i <= 3; i++) {
        // For common cathode RGB LEDs: HIGH turns ON the color, LOW turns it OFF
        bool R = false;  // Default all pins to LOW (OFF)
        bool G = false;
        bool B = false;
        
        // Set the appropriate RGB values based on the traffic light color
        switch (junction.TL_color[i]) {
            case RED:
                R = true;   // Only RED on (HIGH)
                break;
            case YELLOW:
                R = true;   // RED on (HIGH)
                G = true;   // GREEN on (HIGH) - RED + GREEN = YELLOW
                break;
            case GREEN:
                G = true;   // Only GREEN on (HIGH)
                break;
        }
        
        // Apply the RGB values to the corresponding LED
        switch (i) {
            case 1:
                digitalWrite(LED_1_R_PIN, R ? HIGH : LOW);
                digitalWrite(LED_1_G_PIN, G ? HIGH : LOW);
                digitalWrite(LED_1_B_PIN, B ? HIGH : LOW);
                break;
            case 2:
                digitalWrite(LED_2_R_PIN, R ? HIGH : LOW);
                digitalWrite(LED_2_G_PIN, G ? HIGH : LOW);
                digitalWrite(LED_2_B_PIN, B ? HIGH : LOW);
                break;
            case 3:
                digitalWrite(LED_3_R_PIN, R ? HIGH : LOW);
                digitalWrite(LED_3_G_PIN, G ? HIGH : LOW);
                digitalWrite(LED_3_B_PIN, B ? HIGH : LOW);
                break;
        }
    }
}

void update_buttons() {
    unsigned long current_time = millis();
    
    for (int i = 1; i < 4; i++) {
        int pin_num;
        switch (i) {
            case 1: pin_num = BTN_1_PIN; break;
            case 2: pin_num = BTN_2_PIN; break;
            case 3: pin_num = BTN_3_PIN; break;
        }
        
        // Read current button state
        bool current_state = digitalRead(pin_num);
        
        // Default to not pressed
        btn_pressed[i] = false;
        
        // If we detect a potential button press (current HIGH, previous LOW)
        if (current_state && !btn_prev[i]) {
            // Check if enough time has passed since the last press
            if (current_time - last_debounce_time[i] > debounce_delay) {
                // This is a valid button press
                btn_pressed[i] = true;
                last_debounce_time[i] = current_time; // Update last press time
            }
        }
        
        // Update previous state
        btn_prev[i] = current_state;
    }
}

void car_arrivals() {
    unsigned long current_time = millis();
    
    if (mode == MANUAL) {
        // Manual mode: generate cars based on button presses
        if (btn_pressed[1]) {
            // Create new car with unique ID
            int new_id = next_car_id++;
            
            // Add car to lane 1 if not full
            if (lane1.num_of_cars < 4) {
                Car* new_car = new Car(new_id);
                lane1.add_car(new_car);
                Serial.print(current_time);
                Serial.print(", vehicle ");
                Serial.print(new_id);
                Serial.println(", generated");
            } else {
                // Lane is full, car is thrown
                Serial.print(current_time);
                Serial.print(", vehicle ");
                Serial.print(new_id);
                Serial.println(", thrown");
                lane1.thrown_cars++;
                total_thrown_cars++;
                stats.increment_thrown(1);
            }
            Serial.print("Lane 1: ");
            Serial.print(lane1.num_of_cars);
            Serial.print(", Lane 2: ");
            Serial.print(lane2.num_of_cars);
            Serial.print(", Lane 3: ");
            Serial.println(lane3.num_of_cars);
        }
        
        if (btn_pressed[2]) {
            int new_id = next_car_id++;
            
            if (lane2.num_of_cars < 4) {
                Car* new_car = new Car(new_id);
                lane2.add_car(new_car);
                Serial.print(current_time);
                Serial.print(", vehicle ");
                Serial.print(new_id);
                Serial.println(", generated");
            } else {
                // Lane is full, car is thrown
                Serial.print(current_time);
                Serial.print(", vehicle ");
                Serial.print(new_id);
                Serial.println(", thrown");
                lane2.thrown_cars++;
                total_thrown_cars++;
                stats.increment_thrown(2);
            }
            Serial.print("Lane 1: ");
            Serial.print(lane1.num_of_cars);
            Serial.print(", Lane 2: ");
            Serial.print(lane2.num_of_cars);
            Serial.print(", Lane 3: ");
            Serial.println(lane3.num_of_cars);
        }
        
        if (btn_pressed[3]) {
            int new_id = next_car_id++;
            
            if (lane3.num_of_cars < 4) {
                Car* new_car = new Car(new_id);
                lane3.add_car(new_car);
                Serial.print(current_time);
                Serial.print(", vehicle ");
                Serial.print(new_id);
                Serial.println(", generated");
            } else {
                // Lane is full, car is thrown
                Serial.print(current_time);
                Serial.print(", vehicle ");
                Serial.print(new_id);
                Serial.println(", thrown");
                lane3.thrown_cars++;
                total_thrown_cars++;
                stats.increment_thrown(3);
            }
            Serial.print("Lane 1: ");
            Serial.print(lane1.num_of_cars);
            Serial.print(", Lane 2: ");
            Serial.print(lane2.num_of_cars);
            Serial.print(", Lane 3: ");
            Serial.println(lane3.num_of_cars);
        }
    } 
    else if (mode == RANDOM) {
        // Random mode: generate cars based on Markov process for each lane
        for (int lane = 1; lane <= 3; lane++) {
            Lane* target_lane = nullptr;
            switch(lane) {
                case 1: target_lane = &lane1; break;
                case 2: target_lane = &lane2; break;
                case 3: target_lane = &lane3; break;
            }
            
            if (lane_state[lane] == STATE_A) {
                // State A: generate car every 5 seconds
                if (current_time - last_car_time[lane] >= 5000) {
                    // Generate a car with unique ID
                    int new_id = next_car_id++;
                    
                    // Check if lane is full
                    if (target_lane->num_of_cars < 4) {
                        Car* new_car = new Car(new_id);
                        target_lane->add_car(new_car);
                        Serial.print(current_time);
                        Serial.print(", vehicle ");
                        Serial.print(new_id);
                        Serial.println(", generated");
                    } else {
                        // Lane is full, car is thrown
                        Serial.print(current_time);
                        Serial.print(", vehicle ");
                        Serial.print(new_id);
                        Serial.println(", thrown");
                        target_lane->thrown_cars++;
                        total_thrown_cars++;
                        stats.increment_thrown(lane);
                    }
                    Serial.print("Lane 1: ");
                    Serial.print(lane1.num_of_cars);
                    Serial.print(", Lane 2: ");
                    Serial.print(lane2.num_of_cars);
                    Serial.print(", Lane 3: ");
                    Serial.println(lane3.num_of_cars);
                    
                    last_car_time[lane] = current_time;
                    
                    // Transition probability A→B: 0.3
                    if (random(100) < 30) {
                        lane_state[lane] = STATE_B;
                        // Draw waiting time for state B (0.5-5 seconds)
                        waiting_time[lane] = random(500, 5001);
                    }
                }
            } 
            else { // STATE_B
                // State B: wait for the drawn waiting time
                if (current_time - last_car_time[lane] >= waiting_time[lane]) {
                    // Always generate a car in State B, regardless of lane fullness
                    int new_id = next_car_id++;
                    
                    // Check if lane is full
                    if (target_lane->num_of_cars < 4) {
                        Car* new_car = new Car(new_id);
                        target_lane->add_car(new_car);
                        Serial.print(current_time);
                        Serial.print(", vehicle ");
                        Serial.print(new_id);
                        Serial.println(", generated");
                    } else {
                        // Lane is full, car is thrown
                        Serial.print(current_time);
                        Serial.print(", vehicle ");
                        Serial.print(new_id);
                        Serial.println(", thrown");
                        target_lane->thrown_cars++;
                        total_thrown_cars++;
                        stats.increment_thrown(lane);
                    }
                    Serial.print("Lane 1: ");
                    Serial.print(lane1.num_of_cars);
                    Serial.print(", Lane 2: ");
                    Serial.print(lane2.num_of_cars);
                    Serial.print(", Lane 3: ");
                    Serial.println(lane3.num_of_cars);
                    
                    last_car_time[lane] = current_time;
                    
                    // Draw new waiting time
                    waiting_time[lane] = random(500, 5001);
                    
                    // Transition probability B→A: 0.2
                    if (random(100) < 20) {
                        lane_state[lane] = STATE_A;
                    }
                }
            }
        }
    }
    
    // Update statistics
    stats.add_lane_length(1, lane1.num_of_cars);
    stats.add_lane_length(2, lane2.num_of_cars);
    stats.add_lane_length(3, lane3.num_of_cars);
}

void car_departures() {
    unsigned long current_time = millis();
    
    // Only process departures if a green light has been on for at least 1 second
    // and not in a yellow transition
    if (!junction.yellow_transition && (current_time - junction.last_state_change >= 1000) && 
        (current_time - last_departure_time >= 1000)) {
        
        // Determine which lane should have cars departing based on junction state
        Lane* active_lane = nullptr;
        int active_lane_num = 0;
        
        switch (junction.current_state) {
            case GREEN_1:
                active_lane = &lane1;
                active_lane_num = 1;
                break;
            case GREEN_2:
                active_lane = &lane2;
                active_lane_num = 2;
                break;
            case GREEN_3:
                active_lane = &lane3;
                active_lane_num = 3;
                break;
            default:
                // Not in a green state
                return;
        }
        
        // If there's an active lane with cars
        if (active_lane != nullptr && active_lane->num_of_cars > 0) {
            // Remove a car from the lane
            Car* departed_car = active_lane->sub_car();
            
            if (departed_car != nullptr) {
                // Calculate wait time
                unsigned long wait_time = current_time - departed_car->arrival_time;
                stats.add_wait_time(active_lane_num, wait_time);
                
                // Log the junction crossing event
                Serial.print(current_time);
                Serial.print(", vehicle ");
                Serial.print(departed_car->car_ID);
                Serial.println(", crossing junction");
                
                // Log lane occupancy
                Serial.print("Lane 1: ");
                Serial.print(lane1.num_of_cars);
                Serial.print(", Lane 2: ");
                Serial.print(lane2.num_of_cars);
                Serial.print(", Lane 3: ");
                Serial.println(lane3.num_of_cars);
                
                // Free memory
                delete departed_car;
                
                // Update display because car count changed
                display_changed = true;
            }
        }
        
        // Increment green duration counter
        green_duration++;
        
        // Reset timer for next departure
        last_departure_time = current_time;
    }
}

void update_junction_state() {
    unsigned long current_time = millis();
    
    // First check if we're in a yellow transition
    if (junction.yellow_transition) {
        junction.process_yellow_transition();
        return;
    }
    
    // Get the current active traffic light based on junction state
    int current_tl = 0;
    switch (junction.current_state) {
        case GREEN_1: current_tl = 1; break;
        case GREEN_2: current_tl = 2; break;
        case GREEN_3: current_tl = 3; break;
        default: return; // Invalid state or already in transition
    }
    
    // Calculate how long the current light has been green
    unsigned long time_since_change = current_time - junction.last_state_change;
    
    // Check if it's time to change the state
    // Check if the configured green duration has elapsed (in milliseconds)
    if (time_since_change >= junction.green_duration[current_tl] * 1000) {
        // Determine the next state in the sequence
        JunctionState next_state;
        switch (junction.current_state) {
            case GREEN_1:
                next_state = GREEN_2;
                break;
            case GREEN_2:
                next_state = GREEN_3;
                break;
            case GREEN_3:
                next_state = GREEN_1;
                break;
            default:
                return; // Invalid state
        }
        
        // Start the transition to yellow
        junction.start_yellow_transition(next_state);
    }
}

void optimize_traffic_pattern() {
    unsigned long current_time = millis();
    
    // Only run optimization every 30 seconds
    if (current_time - last_optimization_time < OPTIMIZATION_INTERVAL) {
        return;
    }
    
    // Log the optimization event
    Serial.print(current_time);
    Serial.println(", system, Traffic Light Pattern Optimization");
    
    // Calculate scores for each lane
    float lane_scores[4] = {0.0, 0.0, 0.0, 0.0}; // Index 0 unused
    
    // Calculate scores based on average lane length and number of thrown cars
    for (int lane = 1; lane <= 3; lane++) {
        float avg_length = stats.calc_avg_lane_length(lane);
        int thrown = stats.get_thrown(lane);
        float avg_wait = stats.calc_avg_wait_time(lane);
        
        // Score formula: weighted combination of average queue length, thrown cars, and average wait time
        // This gives more weight to thrown cars since they represent lost traffic
        lane_scores[lane] = avg_length * 1.0 + thrown * 2.0 + avg_wait * 0.001;
        
        Serial.print("Lane ");
        Serial.print(lane);
        Serial.print(" score: ");
        Serial.print(lane_scores[lane]);
        Serial.print(" (avg length: ");
        Serial.print(avg_length);
        Serial.print(", thrown: ");
        Serial.print(thrown);
        Serial.print(", avg wait: ");
        Serial.print(avg_wait);
        Serial.println("ms)");
    }
    
    // Find lane with highest score (worst congestion)
    int worst_lane = 1;
    for (int lane = 2; lane <= 3; lane++) {
        if (lane_scores[lane] > lane_scores[worst_lane]) {
            worst_lane = lane;
        }
    }
    
    // Find lane with lowest score (best traffic flow)
    int best_lane = 1;
    for (int lane = 2; lane <= 3; lane++) {
        if (lane_scores[lane] < lane_scores[best_lane]) {
            best_lane = lane;
        }
    }
    
    // Current pattern before changes
    Serial.print("Current pattern: ");
    Serial.print(junction.green_duration[1]);
    Serial.print("-");
    Serial.print(junction.green_duration[2]);
    Serial.print("-");
    Serial.println(junction.green_duration[3]);
    
    // Store original values to check if we need to update
    int original_durations[4];
    for (int i = 1; i <= 3; i++) {
        original_durations[i] = junction.green_duration[i];
    }
    
    // Create new light pattern
    int new_durations[4];
    for (int i = 1; i <= 3; i++) {
        new_durations[i] = junction.green_duration[i];
    }
    
    // Increase green time for worst lane by 1 second (up to max 7 seconds)
    if (new_durations[worst_lane] < 7) {
        new_durations[worst_lane]++;
    }
    
    // Decrease green time for best lane by 1 second (down to min 1 second)
    if (new_durations[best_lane] > 1) {
        new_durations[best_lane]--;
    }
    
    // Check if the pattern actually changed
    bool pattern_changed = false;
    for (int i = 1; i <= 3; i++) {
        if (original_durations[i] != new_durations[i]) {
            pattern_changed = true;
            break;
        }
    }
    
    // Log the result
    if (pattern_changed) {
        Serial.print(current_time);
        Serial.print(", system, Traffic Light Pattern Update: ");
        Serial.print(new_durations[1]);
        Serial.print("-");
        Serial.print(new_durations[2]);
        Serial.print("-");
        Serial.println(new_durations[3]);
        
        // Update the pattern - will be applied after the current round is completed
        junction.update_pattern(new_durations[1], new_durations[2], new_durations[3]);
        
        // Reset thrown car counters after optimization is applied
        total_thrown_cars = 0;
        lane1.thrown_cars = 0;
        lane2.thrown_cars = 0;
        lane3.thrown_cars = 0;
        stats.reset_thrown_cars();
        
        // Update display since pattern changed
        display_changed = true;
    } else {
        Serial.println("No changes to traffic light pattern needed");
    }
    
    // Update the last optimization time
    last_optimization_time = current_time;
}

void setup() {
    setup_lcd();
    setup_serial();
    pin_modes_setup();
    
    // Initialize random seed
    randomSeed(analogRead(A0));
    
    // Set initial display
    display_changed = true;
    
    Serial.print("Smart Junction System initialized at ");
    Serial.println(millis());
}

void loop() {
    // Update inputs and buttons
    update_inputs();
    update_buttons();
    
    // Process traffic light state changes
    update_junction_state();
    
    // Handle car arrivals and departures
    car_departures();
    car_arrivals();
    
    // Update LED states
    update_leds();
    
    // Periodically optimize traffic pattern
    optimize_traffic_pattern();
    
    // Update the LCD display if needed
    if (display_changed) {
        update_monitor();
        display_changed = false; // Reset the flag after updating
    }
    
    // Small delay to prevent CPU hogging
    delay(10);
}