from machine import Pin, PWM, I2C
import time
import ssd1306

i2c = I2C(0, sda=Pin(20), scl=Pin(21), freq=400000)
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

led_rojo = Pin(19, Pin.OUT)
led_ambar = Pin(22, Pin.OUT)
led_verde = Pin(26, Pin.OUT)

boton_start = Pin(15, Pin.IN, Pin.PULL_UP)

def torreta_estado(rojo, ambar, verde):
    led_rojo.value(rojo)
    led_ambar.value(ambar)
    led_verde.value(verde)

pins_servos = [2, 3, 4, 5, 6]
pwms = []

for pin in pins_servos:
    pwm = PWM(Pin(pin))
    pwm.freq(50)
    pwms.append(pwm)

def set_angulo(servo_idx, angulo):
    angulo = max(0, min(180, angulo))
    duty_ns = int(500000 + (angulo / 180) * 2000000)
    pwms[servo_idx].duty_ns(duty_ns)

def mover_suave(servo_idx, angulo_actual, angulo_destino, delay=0.015):
    step = 1 if angulo_destino > angulo_actual else -1
    for a in range(angulo_actual, angulo_destino + step, step):
        set_angulo(servo_idx, a)
        time.sleep(delay)

def actualizar_pantalla(modo, estado_torreta, angulos):
    oled.fill(0)
    oled.text("BRAZO ROBOTICO", 0, 0)
    oled.text(f"Modo: {modo}", 0, 12)
    oled.text(f"Torreta: {estado_torreta}", 0, 24)
    oled.text(f"S1:{angulos[0]} S2:{angulos[1]} S3:{angulos[2]}", 0, 40)
    oled.text(f"S4:{angulos[3]} S5:{angulos[4]}", 0, 52)
    oled.show()

def rutina_manual_secuencial(pos_actual):
    torreta_estado(rojo=0, ambar=1, verde=0)
    
    secuencia_manual = [
        (0, 120, "1. Giro Base"),
        (1, 45,  "2. Bajar Hombro"),
        (2, 135, "3. Extender Codo"),
        (3, 45,  "4. Flex Muñeca"),
        (4, 90,  "5. Cierra Gripper"),
        (4, 0,   "6. Abre Gripper"),
        (0, 90,  "7. Centrar Base")
    ]
    
    for idx_servo, ang_destino, etiqueta in secuencia_manual:
        mover_suave(idx_servo, pos_actual[idx_servo], ang_destino)
        pos_actual[idx_servo] = ang_destino
        actualizar_pantalla(f"MAN:{etiqueta}", "AMBAR", pos_actual)
        time.sleep(0.5)

def rutina_automatica_pallet(pos_actual):
    torreta_estado(rojo=0, ambar=0, verde=1)
    
    secuencia_auto = [
        (0, 30,  "1. Posicion A"),
        (1, 45,  "2. Bajar A"),
        (2, 120, "3. Codo A"),
        (4, 120, "4. Gripper Pick"),
        (1, 90,  "5. Elevar Carga"),
        (0, 150, "6. Giro Pallet B"),
        (1, 45,  "7. Bajar B"),
        (4, 0,   "8. Soltar Carga"),
        (1, 90,  "9. Elevar Vacio"),
        (0, 90,  "10. Home Base")
    ]
    
    for idx_servo, ang_destino, etiqueta in secuencia_auto:
        mover_suave(idx_servo, pos_actual[idx_servo], ang_destino)
        pos_actual[idx_servo] = ang_destino
        actualizar_pantalla(f"AUTO:{etiqueta}", "VERDE", pos_actual)
        time.sleep(0.4)

pos_home = [90, 90, 90, 90, 0]

for i in range(5):
    set_angulo(i, pos_home[i])

while True:
    torreta_estado(rojo=1, ambar=0, verde=0)
    actualizar_pantalla("ESPERANDO...", "ROJO", pos_home)
    
    if boton_start.value() == 0:
        time.sleep_ms(50)
        if boton_start.value() == 0:
            while boton_start.value() == 0:
                time.sleep_ms(10)
            
            pos_actual = list(pos_home)
            rutina_manual_secuencial(pos_actual)
            rutina_automatica_pallet(pos_actual)
            
            for i in range(5):
                mover_suave(i, pos_actual[i], pos_home[i])
                pos_actual[i] = pos_home[i]
    
    time.sleep_ms(100)