from machine import Pin, PWM
import time

# ================= PINS =================
red = Pin(25, Pin.OUT)
yellow = Pin(26, Pin.OUT)
green = Pin(27, Pin.OUT)

buzzer = PWM(Pin(14))
buzzer.duty(0)

switch = Pin(13, Pin.IN, Pin.PULL_UP)

segments = [
    Pin(4, Pin.OUT), Pin(18, Pin.OUT), Pin(21, Pin.OUT),
    Pin(19, Pin.OUT), Pin(22, Pin.OUT), Pin(23, Pin.OUT), Pin(5, Pin.OUT)
]
tens = Pin(15, Pin.OUT)

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),
}

# ================= STATES =================
IDLE, RUNNING = 0, 1
state = IDLE

# ================= DISPLAY =================
def show(n):
    tens.value(1 if n >= 10 else 0)
    u = n - 10 if n >= 10 else n
    for i in range(7):
        segments[i].value(digits[u][i])

def clear_all():
    red.off(); yellow.off(); green.off()
    buzzer.duty(0)
    tens.off()
    for s in segments: s.off()

# ================= SWITCH (DEBOUNCED) =================
def hold_time():
    if switch.value() == 0:
        t0 = time.time()
        while switch.value() == 0:
            time.sleep(0.01)
        return time.time() - t0
    return 0

# ================= BUZZER =================
def beep(freq, sec):
    buzzer.freq(freq)
    buzzer.duty(512)
    time.sleep(sec)
    buzzer.duty(0)

def red_beep():
    buzzer.freq(1000)
    buzzer.duty(512)
    time.sleep(0.3)
    buzzer.duty(0)
    time.sleep(0.7)

def yellow_beep():
    buzzer.freq(1200)
    for _ in range(5):
        buzzer.duty(512)
        time.sleep(0.1)
        buzzer.duty(0)
        time.sleep(0.1)

# ================= STOP CHECK =================
def check_stop():
    if switch.value() == 0:
        t0 = time.time()
        while switch.value() == 0:
            if time.time() - t0 >= 3:
                beep(800, 1)
                clear_all()
                return True
            time.sleep(0.05)
    return False

# ================= TRAFFIC LOOP =================
def traffic_loop():
    while True:

        # RED
        red.on(); yellow.off(); green.off()
        for t in range(15, -1, -1):
            show(t)
            if t > 0:
                red_beep()
            else:
                time.sleep(1)
            if check_stop(): return

        # GREEN
        red.off(); green.on()
        for t in range(10, 4, -1):
            show(t)
            time.sleep(1)
            if check_stop(): return

        # YELLOW
        green.off(); yellow.on()
        for t in range(4, -1, -1):
            show(t)
            if t > 0:
                yellow_beep()
            else:
                time.sleep(1)
            if check_stop(): return

# ================= MAIN =================
clear_all()

while True:
    # IDLE
    if state == IDLE:
        clear_all()
        if switch.value() == 0:
            t = hold_time()
            if t >= 1:
                beep(1000, 2)
                state = RUNNING

    # RUNNING
    elif state == RUNNING:
        traffic_loop()
        state = IDLE