import board
import digitalio
import time

# -----------------------------
# FIXED TIMES
# -----------------------------
y = 3
z = 10

# -----------------------------
# LED SETUP
# -----------------------------
red = digitalio.DigitalInOut(board.GP16)
yellow = digitalio.DigitalInOut(board.GP17)
green = digitalio.DigitalInOut(board.GP18)

for led in [green, yellow, red]:
    led.direction = digitalio.Direction.OUTPUT
    led.value = False

# -----------------------------
# START BUTTON
# -----------------------------
sw1 = digitalio.DigitalInOut(board.GP15)
sw1.direction = digitalio.Direction.INPUT
sw1.pull = digitalio.Pull.UP

# -----------------------------
# MATRIX KEYPAD SETUP
# -----------------------------
rows = [board.GP13, board.GP12, board.GP11, board.GP10]
cols = [board.GP9, board.GP8, board.GP7, board.GP6]

row_pins = []
col_pins = []

for r in rows:
    pin = digitalio.DigitalInOut(r)
    pin.direction = digitalio.Direction.OUTPUT
    pin.value = True
    row_pins.append(pin)

for c in cols:
    pin = digitalio.DigitalInOut(c)
    pin.direction = digitalio.Direction.INPUT
    pin.pull = digitalio.Pull.UP
    col_pins.append(pin)

keys = [
    ["1","2","3","A"],
    ["4","5","6","B"],
    ["7","8","9","C"],
    ["*","0","#","D"]
]

# -----------------------------
# READ MATRIX KEYPAD
# -----------------------------
def read_key():
    for i, r in enumerate(row_pins):
        r.value = False

        for j, c in enumerate(col_pins):
            if not c.value:
                key = keys[i][j]

                # wait release
                while not c.value:
                    pass

                r.value = True
                return key

        r.value = True

    return None

# -----------------------------
# BCD DISPLAY (CD4511)
# -----------------------------
bcd_pins = []
for pin in [board.GP5, board.GP2, board.GP3, board.GP4]:
    p = digitalio.DigitalInOut(pin)
    p.direction = digitalio.Direction.OUTPUT
    bcd_pins.append(p)

# NEW latch pins (IMPORTANT)
s72 = digitalio.DigitalInOut(board.GP1)
s71 = digitalio.DigitalInOut(board.GP0)
s72.direction = digitalio.Direction.OUTPUT
s71.direction = digitalio.Direction.OUTPUT
s72.value = True
s71.value = True

def send_bcd(val):
    for i in range(4):
        bcd_pins[i].value = (val >> i) & 1

def latch(pin):
    pin.value = False
    time.sleep(0.001)
    pin.value = True

def display(num):
    tens = num // 10
    ones = num % 10

    send_bcd(tens)
    latch(s72)

    send_bcd(ones)
    latch(s71)

# -----------------------------
# GET INPUT
# -----------------------------
def get_input():
    value = ""
    print("Enter value (0-99), press #")

    while True:
        key = read_key()

        if key:
            print("Pressed:", key)

            if key.isdigit():
                if len(value) < 2:
                    value += key

            elif key == "#":
                if value != "":
                    return int(value)
                else:
                    print("Error: No input")

            else:
                print("Invalid key")

            time.sleep(0.2)

# -----------------------------
# COUNTDOWN PHASE
# -----------------------------
def run_phase(sec, led, name):
    print(name, "PHASE START")
    led.value = True

    for i in range(sec, -1, -1):
        display(i)
        print(name, ":", i)
        time.sleep(0.7)

    led.value = False

# -----------------------------
# MAIN
# -----------------------------
print("Traffic Light Ready")

# wait for button
while sw1.value:
    pass

x = get_input()

while True:
    run_phase(x, green, "GREEN")
    run_phase(y, yellow, "YELLOW")
    run_phase(z, red, "RED")