from machine import Pin, I2C, freq
from steptime import STEPTIME
from ws2812 import WS2812
from ssd1306 import SSD1306_I2C
import utime
import random

freq(250_000_000)

ROUNDS = 5          
PENALTY_MS = 500 
BEST_FILE = "best.txt"

LOOP = 200     
SETTLE = 20000        
THRESH = 10000      

i2c = I2C(1, scl=Pin(7), sda=Pin(6), freq=400_000)
oled = SSD1306_I2C(128, 64, i2c)

Pin(11, Pin.OUT).value(1)

pixel = WS2812(12, 1, 0.3, 6)

BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
ORANGE = (255, 80, 0)

def set_pixel(color):
    pixel.pixels_fill(color)
    pixel.pixels_show()

PADS = [
    ("Q0", 3,  110, 22),
    ("Q1", 4,   92, 38),
    ("Q2", 2,   20, 46),
    ("Q3", 27,   4, 30),
    ("Q4", 1,   36, 30),
    ("Q5", 26,  20, 14),
]
SIZE = 12  
sensors = []
for i, (name, gpio, x, y) in enumerate(PADS):
    Pin(gpio, Pin.IN, Pin.PULL_UP)
    sensors.append(STEPTIME(i, gpio))   

baseline = [4294967296] * len(PADS) 

##Sensing

def read_pads():

    for sm in sensors:          # 1. start all measurements
        sm.put(LOOP)
        sm.put(SETTLE)
    touched = []
    for i, sm in enumerate(sensors):   # 2. collect the results
        # The PIO counts down from 0, so the raw value wraps around 2**32.
        # Subtracting it from 2**32 gives the positive count.
        value = 4294967296 - sm.get()
        if value < baseline[i]:
            baseline[i] = value
        touched.append(value - baseline[i] > THRESH)
    return touched


def wait_press():
    """Block until any pad pressed"""
    while True:
        touched = read_pads()
        for i in range(len(PADS)):
            if touched[i]:
                return i


def wait_release():
    while any(read_pads()):
        pass


##Display
def centered(text, y):
    oled.text(text, (128 - 8 * len(text)) // 2, y)


def screen(lines):
    oled.fill(0)
    y = (64 - 12 * len(lines)) // 2      # centre the block vertically
    for line in lines:
        centered(line, y)
        y += 12
    oled.show()

def draw_pads(target, label):
    oled.fill(0)
    oled.text(label, 52, 34)
    for i, (name, gpio, x, y) in enumerate(PADS):
        if i == target:
            oled.fill_rect(x, y, SIZE, SIZE, 1)
        else:
            oled.rect(x, y, SIZE, SIZE, 1)
    oled.show()


##File loading/ writing
def load_best():
    try:
        with open(BEST_FILE) as f:
            return int(f.read())
    except (OSError, ValueError):   # file missing or unreadable
        return None


def save_best(value):
    with open(BEST_FILE, "w") as f:
        f.write(str(value))



## Game logic


def play_round():
    ## Waiting
    set_pixel(RED)
    draw_pads(None, "wait")
    delay = random.randint(1000, 3000)
    start = utime.ticks_ms()
    while utime.ticks_diff(utime.ticks_ms(), start) < delay:
        if any(read_pads()):
            return None                  # false start
    ##Go
    target = random.randint(0, len(PADS) - 1)
    set_pixel(GREEN)
    draw_pads(target, "GO!")
    t0 = utime.ticks_ms()                # after oled.show(): target is visible
    penalty = 0
    while True:
        i = wait_press()
        if i == target:
            # ticks_diff handles the wrap around of the millisecond counter
            return utime.ticks_diff(utime.ticks_ms(), t0) + penalty
        penalty += PENALTY_MS            # wrong pad: flash orange, keep going
        set_pixel(ORANGE)
        wait_release()
        set_pixel(GREEN)


def main():
    screen(["REACTION", "", "hands off", "calibrating"])
    for _ in range(30):
        read_pads()

    best = load_best()
    while True:
        # Start screen
        set_pixel(BLUE)
        best_line = "best: %d ms" % best if best else "best: --"
        screen(["REACTION", "", "touch a pad", "to start", best_line])
        wait_press()
        wait_release()

        # Play ROUNDS valid rounds (false starts are repeated)
        times = []
        while len(times) < ROUNDS:
            result = play_round()
            wait_release()
            if result is None:
                set_pixel(RED)
                screen(["Too early!", "", "again..."])
                utime.sleep(1)
                continue
            times.append(result)
            set_pixel(BLACK)
            screen(["%d ms" % result, "", "round %d/%d" % (len(times), ROUNDS)])
            utime.sleep(1)

        # Result screen
        avg = sum(times) // len(times)
        lines = ["average", "%d ms" % avg, "fastest %d ms" % min(times)]
        if best is None or avg < best:
            best = avg
            save_best(best)
            lines.append("NEW RECORD!")
            set_pixel(GREEN)
        else:
            set_pixel(BLUE)
        print("rounds:", times, "average:", avg)   # visible in the terminal
        screen(lines)
        utime.sleep(2)
        wait_press()                     # touch to return to the start screen
        wait_release()



try:
    main()
finally:
    set_pixel(BLACK)
    oled.fill(0)
    oled.show()