from unihiker import GUI
import serial
import serial.tools.list_ports
import time
from pinpong.board import Board, Pin

Board().begin()
btn_success_pin = Pin(Pin.P3, Pin.IN)
btn_fail_pin = Pin(Pin.P4, Pin.IN)

u_gui = GUI()

u_gui.draw_text(text="Test 1", x=70, y=5, font_size=20, color="#0000FF")

status_text = u_gui.draw_text(text="Initializing...", x=50, y=35, font_size=10, color="gray")

u_gui.draw_text(text="Current Step:", x=35, y=55, font_size=15, color="black")
info_text = u_gui.draw_text(text="-- / 11", x=55, y=75, font_size=40, color="#FF0000")

u_gui.draw_text(text="Success(P3)", x=10, y=130, font_size=12, color="green")
success_text = u_gui.draw_text(text="0", x=35, y=150, font_size=25, color="green")

u_gui.draw_text(text="Fail(P4)", x=140, y=130, font_size=12, color="red")
fail_text = u_gui.draw_text(text="0", x=160, y=150, font_size=25, color="red")

u_gui.draw_text(text="Total Cycles:", x=20, y=190, font_size=15, color="black")
cycle_text = u_gui.draw_text(text="0", x=170, y=185, font_size=20, color="blue")

ser = None
ports = list(serial.tools.list_ports.comports())
port_name = None

for p in ports:
    if "ACM" in p.device or "USB" in p.device:
        port_name = p.device
        break

if port_name:
    try:
        ser = serial.Serial(port_name, 9600, timeout=0.1)
        status_text.config(text=f"Conn: {port_name}", color="green")
        info_text.config(text="Ready", color="black", font_size=30)
    except:
        status_text.config(text="Connect Failed", color="red")
else:
    status_text.config(text="No Arduino Found", color="red")

cycle_count = 0
success_count = 0
fail_count = 0
last_p3_state = 0
last_p4_state = 0
waiting_for_judgment = False
stop_stage = 0 

def click_start():
    global stop_stage
    if ser and ser.is_open:
        ser.write(b'1')
        status_text.config(text="Running...", color="blue")
        info_text.config(text="Start", color="blue", font_size=40)
        stop_stage = 0 
        btn_stop.config(text="Stop")

def click_stop():
    global stop_stage, waiting_for_judgment
    if ser and ser.is_open:
        if stop_stage == 0:
            ser.write(b'2') 
            status_text.config(text="Paused", color="orange")
            info_text.config(text="Halt", color="orange", font_size=40)
            btn_stop.config(text="Reset") 
            stop_stage = 1
        else:
            ser.write(b'0')
            status_text.config(text="Resetting...", color="red")
            info_text.config(text="Home", color="green", font_size=40)
            btn_stop.config(text="Stop") 
            stop_stage = 0
            waiting_for_judgment = False

def click_clear():
    global cycle_count, success_count, fail_count
    cycle_count = 0
    cycle_text.config(text="0")
    success_count = 0
    success_text.config(text="0")
    fail_count = 0
    fail_text.config(text="0")
    print("All counts cleared.")

u_gui.add_button(text="Start", x=20, y=230, w=90, h=45, onclick=click_start)
btn_stop = u_gui.add_button(text="Stop", x=130, y=230, w=90, h=45, onclick=click_stop)
u_gui.add_button(text="Clear Data", x=75, y=285, w=90, h=30, onclick=click_clear)

while True:
    current_p3 = btn_success_pin.value()
    current_p4 = btn_fail_pin.value()
    
    if current_p3 == 1 and last_p3_state == 0:
        if waiting_for_judgment:
            success_count += 1
            success_text.config(text=str(success_count))
            print(f"Experiment {cycle_count}: 1")
            waiting_for_judgment = False
        
        if ser and ser.is_open:
            ser.write(b'1')
            info_text.config(text="Next...", color="blue", font_size=40)
            status_text.config(text="Running", color="blue")
                
    if current_p4 == 1 and last_p4_state == 0:
        if waiting_for_judgment:
            fail_count += 1
            fail_text.config(text=str(fail_count))
            print(f"Experiment {cycle_count}: 0")
            waiting_for_judgment = False
            
        if ser and ser.is_open:
            ser.write(b'1')
            info_text.config(text="Next...", color="blue", font_size=40)
            status_text.config(text="Running", color="blue")
        
    last_p3_state = current_p3
    last_p4_state = current_p4

    if ser and ser.is_open and ser.in_waiting > 0:
        try:
            line = ser.readline().decode().strip()
            
            if "STEP:" in line:
                step_num = line.split(":")[1]
                info_text.config(text=f"{step_num} / 11", color="#FF0000", font_size=40)
            
            elif "WAITING" in line:
                info_text.config(text="Wait P3/P4", color="purple", font_size=30)

            elif "CYCLE_DONE" in line:
                cycle_count += 1
                cycle_text.config(text=str(cycle_count))
                waiting_for_judgment = True
                info_text.config(text="Judge?", color="purple", font_size=40)
            
            elif "HOMING" in line:
                 info_text.config(text="Homing...", color="orange", font_size=30)

            elif "STOPPED" in line:
                info_text.config(text="Home", color="green", font_size=35)
                status_text.config(text="Ready", color="gray")
                waiting_for_judgment = False
                
        except:
            pass
            
    time.sleep(0.05)