from machine import Pin
from time import sleep

# Traffic light pins
red = Pin(2, Pin.OUT)
yellow = Pin(3, Pin.OUT)
green = Pin(4, Pin.OUT)

# Pedestrian light pins
ped_red = Pin(5, Pin.OUT)
ped_green = Pin(6, Pin.OUT)

# Pedestrian button (internal pull-up)
button = Pin(28, Pin.IN, Pin.PULL_UP)

# Initial states
def init_lights():
    red.on()
    yellow.off()
    green.off()
    ped_red.on()
    ped_green.off()

def pedestrian_crossing():
    # Stop traffic
    green.off()
    yellow.on()
    sleep(1)
    yellow.off()
    red.on()
    sleep(1)

    # Let pedestrians walk
    ped_red.off()
    ped_green.on()
    print("Pedestrians can walk")
    sleep(5)

    # Flash green pedestrian light
    for _ in range(3):
        ped_green.toggle()
        sleep(0.5)

    ped_green.off()
    ped_red.on()
    print("Pedestrians stop")

    # Resume traffic
    red.off()
    yellow.on()
    sleep(1)
    yellow.off()
    green.on()

def normal_traffic_cycle():
    red.off()
    yellow.off()
    green.on()
    sleep(1)

while True:
    init_lights()
    normal_traffic_cycle()

    if not button.value():  # Button pressed (active LOW)
        print("Pedestrian button pressed")
        pedestrian_crossing()

    sleep(0.1)  # Polling delay