import board
import digitalio
import time

# Define pin assignments
traffic_red_pin = board.GP0
traffic_yellow_pin = board.GP1
traffic_green_pin = board.GP2
pedestrian_red_pin = board.GP3
pedestrian_green_pin = board.GP4
button_pin = board.GP28

# Define time intervals
yellow_duration = 1
red_duration = 2

# Initialize LEDs and button
traffic_red = digitalio.DigitalInOut(traffic_red_pin)
traffic_yellow = digitalio.DigitalInOut(traffic_yellow_pin)
traffic_green = digitalio.DigitalInOut(traffic_green_pin)
pedestrian_green = digitalio.DigitalInOut(pedestrian_green_pin)
pedestrian_red = digitalio.DigitalInOut(pedestrian_red_pin)
button = digitalio.DigitalInOut(button_pin)

for light in [traffic_red, traffic_yellow, traffic_green, pedestrian_green, pedestrian_red]:
    light.direction = digitalio.Direction.OUTPUT

button.direction = digitalio.Direction.INPUT
button.pull = digitalio.Pull.DOWN


# Initialize states
traffic_red.value = False  # Turn off red
traffic_yellow.value = False  # Turn off yellow
traffic_green.value = True  # Turn on green
pedestrian_red.value = True  # Turn on pedestrian red
pedestrian_green.value = False  # Turn off pedestrian green
current_state = "traffic_green"
start_time = time.monotonic()

while True:

    # Update state based on time
    elapsed_time = time.monotonic() - start_time

    if current_state == "traffic_green":
        # Handle button press
        if button.value:
            traffic_red.value = False  # Turn off red
            traffic_yellow.value = True  # Turn off yellow
            traffic_green.value = False  # Turn on green
            pedestrian_red.value = True  # Turn on pedestrian red
            pedestrian_green.value = False  # Turn off pedestrian green
            current_state = "traffic_yellow"
            start_time = time.monotonic()

    elif current_state == "traffic_yellow":
        if elapsed_time >= yellow_duration:
            traffic_red.value = True  # Turn on red
            traffic_yellow.value = False  # Turn off yellow
            traffic_green.value = False  # Turn off green
            pedestrian_red.value = False  # Turn off pedestrian red
            pedestrian_green.value = True  # Turn on pedestrian green
            current_state = "traffic_red"
            start_time = time.monotonic()

    elif current_state == "traffic_red":
        if elapsed_time >= red_duration:
            traffic_red.value = True  # Turn on red
            traffic_yellow.value = True  # Turn on yellow
            traffic_green.value = False  # Turn off green
            pedestrian_red.value = True  # Turn on pedestrian red
            pedestrian_green.value = False  # Turn off pedestrian green
            current_state = "traffic_red_yellow"
            start_time = time.monotonic()

    elif current_state == "traffic_red_yellow":
        if elapsed_time >= yellow_duration:
            traffic_red.value = False  # Turn off red
            traffic_yellow.value = False  # Turn off yellow
            traffic_green.value = True  # Turn on green
            pedestrian_red.value = True  # Turn on pedestrian red
            pedestrian_green.value = False  # Turn off pedestrian green
            current_state = "traffic_green"
            start_time = time.monotonic()