from machine import Pin, I2C
import time
import ssd1306

# ---------- TIMINGS ----------
GREEN_TIME = 20   # green light duration in seconds
YELLOW_TIME = 5   # yellow light duration in seconds

# ---------- OLED SETUP ----------
i2c = I2C(0, scl=Pin(1), sda=Pin(0))  # adjust pins if needed
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

def oled_show(junction, state, t):
    oled.fill(0)
    oled.text("JUNCTION: {}".format(junction), 0, 0)
    oled.text("SIGNAL: {}".format(state), 0, 25)
    oled.text("TIME: {}s".format(t), 0, 45)
    oled.show()

# ---------- LED SETUP ----------
# Each junction has one RGB LED
J = [
    {"R": Pin(2, Pin.OUT),  "G": Pin(4, Pin.OUT),  "B": Pin(3, Pin.OUT)},   # junction 1
    {"R": Pin(5, Pin.OUT),  "G": Pin(7, Pin.OUT),  "B": Pin(6, Pin.OUT)},   # junction 2
    {"R": Pin(8, Pin.OUT),  "G": Pin(10, Pin.OUT), "B": Pin(9, Pin.OUT)},   # junction 3
    {"R": Pin(11, Pin.OUT), "G": Pin(13, Pin.OUT), "B": Pin(12, Pin.OUT)}   # junction 4
]

# ---------- HELPER FUNCTIONS ----------
def all_red():
    """Set all junctions to RED"""
    for j in J:
        j["R"].value(1)
        j["G"].value(0)
        j["B"].value(0)

def run_junction(index):
    """Run one junction through GREEN -> YELLOW -> RED"""
    # ----- GREEN -----
    all_red()                  # other junctions RED
    J[index]["R"].value(0)     # turn off red
    J[index]["G"].value(1)     # turn on green
    for t in range(GREEN_TIME, 0, -1):
        oled_show(index + 1, "GREEN", t)
        time.sleep(1)

    # ----- YELLOW -----
    J[index]["G"].value(0)     # turn off green
    J[index]["R"].value(1)     # red ON
    J[index]["G"].value(1)     # green ON => yellow
    for t in range(YELLOW_TIME, 0, -1):
        oled_show(index + 1, "YELLOW", t)
        time.sleep(1)

    # ----- RED -----
    J[index]["R"].value(1)
    J[index]["G"].value(0)

# ---------- MAIN LOOP ----------
while True:
    for i in range(4):
        run_junction(i)