from machine import Pin
from time import sleep

# LEDs
red = Pin(0, Pin.OUT)
yellow = Pin(1, Pin.OUT)
green = Pin(2, Pin.OUT)

# 7-segment pins (a, b, c, d, e, f, g)
segments = [
    Pin(3, Pin.OUT),  # a
    Pin(4, Pin.OUT),  # b
    Pin(5, Pin.OUT),  # c
    Pin(6, Pin.OUT),  # d
    Pin(7, Pin.OUT),  # e
    Pin(8, Pin.OUT),  # f
    Pin(9, Pin.OUT)   # g
]

# Numbers for common cathode (1=ON, 0=OFF)
digits = {
    0: [1,1,1,1,1,1,0],
    1: [0,1,1,0,0,0,0],
    2: [1,1,0,1,1,0,1],
    3: [1,1,1,1,0,0,1],
    4: [0,1,1,0,0,1,1],
    5: [1,0,1,1,0,1,1],
    6: [1,0,1,1,1,1,1],
    7: [1,1,1,0,0,0,0],
    8: [1,1,1,1,1,1,1],
    9: [1,1,1,1,0,1,1]
}

def display_number(n):
    pattern = digits.get(n, [0,0,0,0,0,0,0])
    for pin, val in zip(segments, pattern):
        pin.value(val)

def countdown(seconds, led_color):
    for sec in range(seconds, 0, -1):
        display_number(sec)
        sleep(1)

def traffic_cycle():
    # RED
    red.value(1)
    yellow.value(0)
    green.value(0)
    countdown(5, "RED")
    red.value(0)

    # GREEN
    red.value(0)
    yellow.value(0)
    green.value(1)
    countdown(5, "GREEN")
    green.value(0)

    # YELLOW
    red.value(0)
    yellow.value(1)
    green.value(0)
    countdown(3, "YELLOW")
    yellow.value(0)

while True:
    traffic_cycle()