#include <LiquidCrystal.h>

//pin constants
const int SWITCH_PIN = 21;
const int BUTTON_PINS[3] = {27, 26, 22}; // Vehicle buttons for lanes 1, 2, 3

// RGB LED pins for traffic lights (3 traffic lights, each with R,G,B)
const int TL_PINS[3][3] = {
  {0, 1, 2},    // TL1: R,G,B
  {3, 4, 5},    // TL2: R,G,B
  {6, 13, 14}   // TL3: R,G,B
};

// Time constants (in milliseconds)
const unsigned long YELLOW_DURATION = 500;    // Yellow light duration: 0.5 seconds
const unsigned long MIN_GREEN_DURATION = 1000; // Min green light: 1 second
const unsigned long MAX_GREEN_DURATION = 7000; // Max green light: 7 seconds
const unsigned long OPTIMIZATION_INTERVAL = 30000; // Optimization interval: 30 seconds

//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
};

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

enum GenerationState {
    CONSTANT,
    VARIABLE
};

// Forward declarations
void log_event(String event);
void log_lanes_status();

// Global variables for generating unique vehicle IDs
int next_vehicle_id = 1;

// Global variables for timing
unsigned long last_optimization_time = 0;
unsigned long current_state_start_time = 0;
unsigned long last_log_time = 0;

// Global variables for statistics
int thrown_cars = 0;
int thrown_cars_current_pattern = 0;

// Variables for random generation
GenerationState lane_states[3] = {CONSTANT, CONSTANT, CONSTANT};
unsigned long lane_next_generation[3] = {0, 0, 0};
float lane_probability[3][2][2] = {
  {{0.7, 0.3}, {0.2, 0.8}},  // Lane 1
  {{0.7, 0.3}, {0.2, 0.8}},  // Lane 2
  {{0.7, 0.3}, {0.2, 0.8}}   // Lane 3
};

// Light pattern durations (in seconds)
int tl_green_duration[3] = {3, 3, 3};

//global variables
Mode mode;
bool display_changed = true;
bool round_completed = false;
unsigned long round_start_time = 0;
bool pattern_optimization_needed = false;

//structs and functions
struct Car;
struct Car {
    int car_ID;     // Unique ID for each car
    unsigned long arrival_time; // Time when the car arrived
    Car* next;      // Pointer to next car in lane
    
    // Constructor
    Car(int id) : car_ID(id), arrival_time(millis()), next(nullptr) {}
};

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
    unsigned long total_wait_time; // For statistics
    unsigned long max_wait_time;   // For statistics
    int total_cars_served;         // For statistics
    int cars_thrown;               // Cars rejected because lane was full
    
    //constructor
    Lane() : head(nullptr), tail(nullptr), num_of_cars(0), 
             total_wait_time(0), max_wait_time(0), total_cars_served(0), cars_thrown(0) {}
    
    // Function to add a car to the end of the linked list (FIFO)
    bool add_car(Car* car) {
        if (car == nullptr) return false;
        
        // Check if lane is full (max 4 cars)
        if (num_of_cars >= 4) {
            cars_thrown++;
            thrown_cars++;
            thrown_cars_current_pattern++;
            log_event(String(millis()) + ", vehicle " + String(car->car_ID) + ", thrown");
            delete car; // Clean up memory
            return false;
        }
        
        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++;
        log_event(String(millis()) + ", vehicle " + String(car->car_ID) + ", generated");
        log_lanes_status();
        return true;
    }
    
    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--;
        
        // Update statistics
        unsigned long wait_time = millis() - temp->arrival_time;
        total_wait_time += wait_time;
        if (wait_time > max_wait_time) {
            max_wait_time = wait_time;
        }
        total_cars_served++;
        
        // Log crossing event
        log_event(String(millis()) + ", vehicle " + String(temp->car_ID) + ", crossing junction");
        log_lanes_status();
        
        return temp;
    }
};

struct TrafficLights {
    JunctionState current_state; // Current state of the junction
    Color TL_colors[3];          // Colors for each traffic light
    
    // Constructor
    TrafficLights() : current_state(GREEN_1) {
        TL_colors[0] = GREEN;  // TL1 starts green
        TL_colors[1] = RED;    // TL2 starts red
        TL_colors[2] = RED;    // TL3 starts red
    }
    
    // Function to change the state with a yellow transition
    void change_state(JunctionState new_state) {
        // Log the traffic light change event
        log_event(String(millis()) + ", system, traffic light changing");
        
        // Determine the next colors based on the new state
        Color next_colors[3] = {RED, RED, RED}; // Default all to red
        
        // Set the green light based on new state
        switch (new_state) {
            case GREEN_1:
                next_colors[0] = GREEN;
                break;
            case GREEN_2:
                next_colors[1] = GREEN;
                break;
            case GREEN_3:
                next_colors[2] = GREEN;
                break;
        }
        
        // Apply yellow transition where needed and check if any yellow transitions are needed
        bool yellow_needed = false;
        for (int i = 0; i < 3; i++) {
            if (TL_colors[i] == GREEN && next_colors[i] == RED) {
                TL_colors[i] = YELLOW;
                yellow_needed = true;
            }
        }
        
        // Update physical lights for yellow state
        if (yellow_needed) {
            update_physical_lights();
            delay(YELLOW_DURATION);
        }
        
        // Set all lights to their final colors
        for (int i = 0; i < 3; i++) {
            TL_colors[i] = next_colors[i];
        }
        
        // Update the current state
        current_state = new_state;
        current_state_start_time = millis();
        
        // Update physical lights
        update_physical_lights();
        
        // Update display since lights changed
        display_changed = true;
    }
    
    // Method to update the physical RGB LEDs based on traffic light states
    void update_physical_lights() {
        // Debug logging
        Serial.print("Updating lights - TL1: ");
        Serial.print(TL_colors[0]);
        Serial.print(", TL2: ");
        Serial.print(TL_colors[1]);
        Serial.print(", TL3: ");
        Serial.println(TL_colors[2]);
        
        // Update all traffic lights
        for (int tl = 0; tl < 3; tl++) {
            switch (TL_colors[tl]) {
                case RED:
                    analogWrite(TL_PINS[tl][0], 1024);  // RED ON
                    analogWrite(TL_PINS[tl][1], 0);   // GREEN OFF
                    digitalWrite(TL_PINS[tl][2], 0);   // BLUE OFF
                    break;
                case YELLOW:
                    analogWrite(TL_PINS[tl][0], 1024);  // RED ON
                    analogWrite(TL_PINS[tl][1], 1024);  // GREEN ON (RED+GREEN=YELLOW)
                    analogWrite(TL_PINS[tl][2], 0);   // BLUE OFF
                    break;
                case GREEN:
                    analogWrite(TL_PINS[tl][0], 0);   // RED OFF
                    analogWrite(TL_PINS[tl][1], 1024);  // GREEN ON
                    analogWrite(TL_PINS[tl][2], 0);   // BLUE OFF
                    break;
            }
        }
    }
};

//initialize structures
Lane lanes[3]; // Array of lanes
TrafficLights junction;

// Function to log events to Serial and file
void log_event(String event) {
    Serial.println(event);
    
    // Note: In a real implementation, this would also write to file
    // But wokwi doesn't support file operations, as mentioned in the comments
    /* 
    File logFile = SD.open("my_junction.log", FILE_WRITE);
    if (logFile) {
        logFile.println(event);
        logFile.close();
    }
    */
}

// Function to log the current status of all lanes
void log_lanes_status() {
    String status = "Lane 1: " + String(lanes[0].num_of_cars) + 
                    ", Lane 2: " + String(lanes[1].num_of_cars) + 
                    ", Lane 3: " + String(lanes[2].num_of_cars);
    Serial.println(status);
}

//setup functions
void setup_lcd() {
    lcd.begin(20, 4);
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Traffic System");
    lcd.setCursor(0, 1);
    lcd.print("Initializing...");
    lcd.clear();
}

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

void setup_inputs() {
    pinMode(SWITCH_PIN, INPUT_PULLUP);
    for (int i = 0; i < 3; i++) {
        pinMode(BUTTON_PINS[i], INPUT);
    }
}

void setup_outputs() {
    // Initialize all traffic light pins
    for (int tl = 0; tl < 3; tl++) {
        for (int color = 0; color < 3; color++) {
            pinMode(TL_PINS[tl][color], OUTPUT);
        }
        
        // Set initial state to RED (for common cathode LEDs)
        digitalWrite(TL_PINS[tl][0], HIGH);  // RED ON
        digitalWrite(TL_PINS[tl][1], LOW);   // GREEN OFF
        digitalWrite(TL_PINS[tl][2], LOW);   // BLUE OFF
    }
    
    // Log initialization
    Serial.println("Traffic lights initialized with RED state");
}

// Helper function to get light state as character
char get_light_char(Color color) {
    switch(color) {
        case GREEN: return 'G';
        case YELLOW: return 'Y';
        case RED: return 'R';
        default: return '?';
    }
}

// Update LCD with current status
void update_monitor() {
    //lcd.clear();
    
    // First row: Mode and Light Pattern
    lcd.setCursor(0, 0);
    lcd.print("Mode: ");
    lcd.print(mode == MANUAL ? "Manual " : "Random ");
    
    // Show all three traffic light states
    lcd.setCursor(0, 1);
    for (int i = 0; i < 3; i++) {
        lcd.print("TL");
        lcd.print(i+1);
        lcd.print(": ");
        lcd.print(get_light_char(junction.TL_colors[i]));
        lcd.print(" ");
    }
    
    // Show light timing pattern
       lcd.setCursor(0, 2);
 
    lcd.print("Pattern: ");
    lcd.print(String(tl_green_duration[0]) + "-" + String(tl_green_duration[1]) + "-" + String(tl_green_duration[2]));
    
    // Second row: Lanes occupancy and thrown cars
    lcd.setCursor(0, 3);
    lcd.print("L:" + String(lanes[0].num_of_cars) + 
              "," + String(lanes[1].num_of_cars) + 
              "," + String(lanes[2].num_of_cars) + 
              " T:" + String(thrown_cars_current_pattern));
              
    display_changed = false;
}

// Handle button inputs for manual car generation
void check_buttons() {
    static bool button_prev[3] = {HIGH, HIGH, HIGH};
    bool button_curr[3];
    
    for (int i = 0; i < 3; i++) {
        button_curr[i] = digitalRead(BUTTON_PINS[i]);
        
        // Check for button press (falling edge)
        if (button_prev[i] == HIGH && button_curr[i] == LOW) {
            Car* new_car = new Car(next_vehicle_id++);
            lanes[i].add_car(new_car);
            display_changed = true;
        }
        
        button_prev[i] = button_curr[i];
    }
}

// Generate cars randomly according to the Markov process
void generate_random_cars() {
    unsigned long current_time = millis();
    
    // Check each lane for car generation
    for (int lane = 0; lane < 3; lane++) {
        if (current_time >= lane_next_generation[lane]) {
            if (lane_states[lane] == CONSTANT) {
                // Generate car every 5 seconds
                Car* new_car = new Car(next_vehicle_id++);
                lanes[lane].add_car(new_car);
                display_changed = true;
                
                lane_next_generation[lane] = current_time + 5000; // 5 seconds
                
                // Check if state should change
                if (random(100) < lane_probability[lane][0][1] * 100) {
                    lane_states[lane] = VARIABLE;
                }
            } else { // VARIABLE state
                // Generate with variable wait time (0.5-5 seconds)
                Car* new_car = new Car(next_vehicle_id++);
                lanes[lane].add_car(new_car);
                display_changed = true;
                
                // Draw waiting time from uniform distribution (0.5-5 seconds)
                unsigned long wait_time = random(500, 5001);
                lane_next_generation[lane] = current_time + wait_time;
                
                // Check if state should change
                if (random(100) < lane_probability[lane][1][0] * 100) {
                    lane_states[lane] = CONSTANT;
                }
            }
        }
    }
}

// Forward declaration for optimize_light_pattern
void optimize_light_pattern();

// Let cars pass through junction when their lane has a green light
void process_junction() {
    unsigned long current_time = millis();
    bool time_expired = false;
    unsigned long time_in_state = current_time - current_state_start_time;
    int current_lane = static_cast<int>(junction.current_state); // Convert enum to int (0, 1, or 2)
    
    // Determine if we should change the state based on time
    if (time_in_state >= tl_green_duration[current_lane] * 1000 || time_in_state >= MAX_GREEN_DURATION) {
        time_expired = true;
    }
    
    // If there's a car waiting in the current lane, let it cross
    static unsigned long car_pass_start_time = 0;
    static bool car_passing = false;
    
    if (car_passing) {
        // Check if 1 second has passed since car started crossing
        if (current_time - car_pass_start_time >= 1000) {
            car_passing = false;
            display_changed = true;
        }
    } else if (lanes[current_lane].head != nullptr) {
        Car* car = lanes[current_lane].sub_car();
        delete car; // Clean up memory
        car_passing = true;
        car_pass_start_time = current_time;
        display_changed = true;
        
        // Reset timer to ensure minimum green time
        if (time_in_state < MIN_GREEN_DURATION) {
            current_state_start_time = current_time - MIN_GREEN_DURATION;
        }
    }
    
    // Check if this is the last lane to complete a round
    if (current_lane == 2 && time_expired) {
        round_completed = true;
    }
    
    // Change to the next state if time expired
    if (time_expired && !car_passing) {
        JunctionState next_state = static_cast<JunctionState>((current_lane + 1) % 3);
        junction.change_state(next_state);
        display_changed = true;
        
        // If we completed a round and optimization is needed, optimize now
        if (round_completed && pattern_optimization_needed) {
            optimize_light_pattern();
            round_completed = false;
            pattern_optimization_needed = false;
        }
    }
}

// Optimize the traffic light pattern based on lane statistics
void optimize_light_pattern() {
    log_event(String(millis()) + ", system, traffic light pattern optimization");
    
    // Calculate lane weights
    float lane_weights[3];
    float total_weight = 0;
    
    for (int i = 0; i < 3; i++) {
        lane_weights[i] = lanes[i].num_of_cars + (lanes[i].cars_thrown * 2);
        total_weight += lane_weights[i];
    }
    
    // Default total weight if all lanes empty
    if (total_weight == 0) total_weight = 3;
    
    // Store old pattern for logging
    String old_pattern = String(tl_green_duration[0]) + "-" + 
                         String(tl_green_duration[1]) + "-" + 
                         String(tl_green_duration[2]);
    
    // Calculate new durations
    for (int i = 0; i < 3; i++) {
        int new_duration = max(1, (int)round(lane_weights[i] / total_weight * 9));
        tl_green_duration[i] = min(7, new_duration); // Cap at maximum 7 seconds
    }
    
    // Log the pattern update
    String new_pattern = String(tl_green_duration[0]) + "-" + 
                         String(tl_green_duration[1]) + "-" + 
                         String(tl_green_duration[2]);
    log_event(String(millis()) + ", system, traffic light pattern update: " + old_pattern + " -> " + new_pattern);
    
    // Reset thrown cars counter for new pattern
    thrown_cars_current_pattern = 0;
    
    // Update display
    display_changed = true;
}

//loop functions
void update_inputs() {
    // Read switch state
    bool switch_state = digitalRead(SWITCH_PIN);
    Mode new_mode = switch_state ? MANUAL : RANDOM;
    
    // If mode changed, update display
    if (mode != new_mode) {
        mode = new_mode;
        display_changed = true;
    }
}

void check_optimization_time() {
    unsigned long current_time = millis();
    
    // Check if it's time to optimize (every 30 seconds)
    if (current_time - last_optimization_time >= OPTIMIZATION_INTERVAL) {
        pattern_optimization_needed = true;
        last_optimization_time = current_time;
    }
}

void setup() {
    setup_lcd();
    setup_serial();
    setup_inputs();
    setup_outputs();
    
    // Initialize random seed
    randomSeed(analogRead(0));
    
    // Initialize timers
    last_optimization_time = millis();
    current_state_start_time = millis();
    
    // Set initial mode
    mode = digitalRead(SWITCH_PIN) ? MANUAL : RANDOM;
    
    // Initialize lane generation times for random mode
    for (int i = 0; i < 3; i++) {
        lane_next_generation[i] = millis() + 5000;
    }
    
    // Start with junction in GREEN_1 state
    junction.change_state(GREEN_1);
    
    display_changed = true;
    
    // Log system start
    log_event(String(millis()) + ", system, traffic control system started");
}

void loop() {
    // Update mode based on switch position
    update_inputs();
    
    // Generate cars based on mode
    if (mode == MANUAL) {
        check_buttons();
    } else { // RANDOM mode
        generate_random_cars();
    }
    
    // Process junction (move cars when they have green light)
    process_junction();
    
    // Check if it's time for optimization
    check_optimization_time();
    
    // Update the LCD display if anything changed
    if (display_changed) {
        update_monitor();
    }
    
    // Small delay to prevent CPU hogging
    delay(50);
}