from machine import Pin, PWM
import utime

from display import Display
from joystick import Joystick
from led_matrix import LedMatrix
from buzzer import Buzzer

from envelope import Envelope
from control_law import ControlLaw, MODE_DIRECT, MODE_NORMAL
from cas_logic import CasLogic, SEV_INFO, SEV_CAUTION, SEV_WARNING

ticks_ms = utime.ticks_ms
ticks_add = utime.ticks_add
ticks_diff = utime.ticks_diff

# Pinagem (BitDogLab v7 espelho)
PIN_BTN_A = 5
PIN_BUZZ = 21
PIN_NP = 7
PIN_JOY_X = 27
PIN_JOY_Y = 26
PIN_JOY_SW = 22
PIN_I2C_SDA = 2
PIN_I2C_SCL = 3
OLED_ADDR = 0x3C

# LED RGB (se estiver no seu diagrama; se não, pode remover)
PIN_LED_R = 13
PIN_LED_G = 11
PIN_LED_B = 12

_btn_event = False
def _btn_irq(_pin):
    global _btn_event
    _btn_event = True

def _pwm_init(pin, freq=1000):
    p = PWM(Pin(pin, Pin.OUT))
    p.freq(freq)
    p.duty_u16(0)
    return p

def _rgb_set(p_r, p_g, p_b, r, g, b):
    p_r.duty_u16(int(r)); p_g.duty_u16(int(g)); p_b.duty_u16(int(b))

def _blink(now_ms, period_ms):
    return ((int(now_ms) // int(period_ms)) & 1) == 0

def _rgb_apply(p_r, p_g, p_b, mode, cas_code, envelope_active, now_ms):
    code = str(cas_code).upper()
    m = str(mode).upper()

    if code == "WARNING":
        on = _blink(now_ms, 150)
        _rgb_set(p_r, p_g, p_b, 25000 if on else 0, 0, 0)
        return
    if envelope_active or code == "ENVELOPE":
        on = _blink(now_ms, 300)
        _rgb_set(p_r, p_g, p_b, 18000 if on else 0, 6000 if on else 0, 0)
        return
    if m == "DIRECT":
        _rgb_set(p_r, p_g, p_b, 0, 0, 20000)
    else:
        _rgb_set(p_r, p_g, p_b, 0, 20000, 0)

def _pick_worst(cas_x, cas_y):
    # retorna o pior (maior severity). se empate, prioriza WARNING>ENVELOPE>LIMIT>OVR>outros
    def rank(c):
        sev = int(c.get("severity", 0))
        code = str(c.get("code", "")).upper()
        bonus = 0
        if code == "WARNING": bonus = 50
        elif code == "ENVELOPE": bonus = 40
        elif code == "LIMIT": bonus = 30
        elif code == "OVR": bonus = 20
        return sev * 100 + bonus
    return cas_x if rank(cas_x) >= rank(cas_y) else cas_y

# IRQ do botão A
btn = Pin(PIN_BTN_A, Pin.IN, Pin.PULL_UP)
btn.irq(trigger=Pin.IRQ_FALLING, handler=_btn_irq)

# Drivers
disp = Display(i2c_id=1, sda_pin=PIN_I2C_SDA, scl_pin=PIN_I2C_SCL, addr=OLED_ADDR)
print("I2C scan:", [hex(a) for a in disp.scan()])
disp.splash("BitDogLab v7", "mini_fbw XY")

# Joystick (ajuste invert_y se o seu Y estiver “ao contrário”)
js = Joystick(pin_x=PIN_JOY_X, pin_y=PIN_JOY_Y, pin_sw=PIN_JOY_SW, dead_u16=1500, invert_y=True)
js.start_calibration(now_ms=ticks_ms(), duration_ms=600, sample_period_ms=10)

mx = LedMatrix(pin=PIN_NP, serpentine=True, brightness=0.12)
bz = Buzzer(pin=PIN_BUZZ)

pwm_r = _pwm_init(PIN_LED_R)
pwm_g = _pwm_init(PIN_LED_G)
pwm_b = _pwm_init(PIN_LED_B)

# Core: dois canais independentes
env_x = Envelope(limit_pct=80.0, max_cmd_threshold_pct=95.0, max_cmd_hold_ms=3000, max_cmd_hysteresis_pct=2.0)
env_y = Envelope(limit_pct=80.0, max_cmd_threshold_pct=95.0, max_cmd_hold_ms=3000, max_cmd_hysteresis_pct=2.0)

cl_x = ControlLaw(normal_limit_pct=80.0, alpha_fixed=0.25, tau_ms=None)
cl_y = ControlLaw(normal_limit_pct=80.0, alpha_fixed=0.25, tau_ms=None)

cas = CasLogic(limit_pct=80.0, ovr_threshold_pct=95.0, warning_hold_ms=3000)

mode = MODE_DIRECT
cl_x.set_mode(mode, warm_start_cmd_pct=0.0)
cl_y.set_mode(mode, warm_start_cmd_pct=0.0)

# Debounce
DEBOUNCE_MS = 200
last_toggle_ms = 0

# Scheduler
CTRL_MS = 20
OLED_MS = 100
MATRIX_MS = 50
RGB_MS = 50
BUZ_MS = 10

t_ctrl = ticks_ms()
t_oled = t_ctrl
t_mat = t_ctrl
t_rgb = t_ctrl
t_buz = t_ctrl

# Estado
cmd_x = 0.0
cmd_y = 0.0
act_x = 0.0
act_y = 0.0
lim_x = 0.0
lim_y = 0.0
env_act_x = False
env_act_y = False
alpha_x = 0.0
alpha_y = 0.0

cas_out = {"code": "DIRECT", "text": "DIRECT", "invert": False, "blink": False, "severity": SEV_INFO}
prev_warn_overall = False

while True:
    now = ticks_ms()

    # Botão A: troca modo
    if _btn_event:
        _btn_event = False
        if ticks_diff(now, last_toggle_ms) > DEBOUNCE_MS:
            last_toggle_ms = now
            mode = MODE_NORMAL if mode == MODE_DIRECT else MODE_DIRECT
            cl_x.set_mode(mode, warm_start_cmd_pct=cmd_x)
            cl_y.set_mode(mode, warm_start_cmd_pct=cmd_y)
            bz.beep(now, ms=80, freq=2200)
            print("Mode:", mode)

    # Controle (20ms)
    if ticks_diff(now, t_ctrl) >= 0:
        t_ctrl = ticks_add(t_ctrl, CTRL_MS)

        js.update(now_ms=now)
        x, y, sw = js.read_pct()
        cmd_x = x
        cmd_y = y

        act_x, lim_x, env_act_x, alpha_x = cl_x.step(cmd_x, mode=mode, now_ms=now)
        act_y, lim_y, env_act_y, alpha_y = cl_y.step(cmd_y, mode=mode, now_ms=now)

        hold_x, max_active_x, trig_x = env_x.update_max_cmd(cmd_x, now_ms=now)
        hold_y, max_active_y, trig_y = env_y.update_max_cmd(cmd_y, now_ms=now)

        # CAS por eixo (mesma classe, avaliando separadamente)
        cas_x = cas.evaluate(mode=mode, cmd_raw_pct=cmd_x, cmd_limited_pct=lim_x,
                             envelope_active=env_act_x, max_cmd_hold_ms=hold_x, max_cmd_triggered=trig_x)
        cas_y = cas.evaluate(mode=mode, cmd_raw_pct=cmd_y, cmd_limited_pct=lim_y,
                             envelope_active=env_act_y, max_cmd_hold_ms=hold_y, max_cmd_triggered=trig_y)

        worst = _pick_worst(cas_x, cas_y)

        # Marcar qual eixo disparou (X/Y) no texto
        axis = "X" if worst is cas_x else "Y"
        cas_out = {
            "code": worst.get("code", ""),
            "severity": worst.get("severity", SEV_INFO),
            "invert": worst.get("invert", False),
            "blink": worst.get("blink", False),
            "text": "%s %s" % (worst.get("code", ""), axis),
        }

        # WARNING geral se qualquer eixo triggar
        warn_overall = (str(cas_out["code"]).upper() == "WARNING")
        if warn_overall and (not prev_warn_overall):
            bz.trigger_warning_3s(now, freq=2400)
        prev_warn_overall = warn_overall

    # Buzzer update
    if ticks_diff(now, t_buz) >= 0:
        t_buz = ticks_add(t_buz, BUZ_MS)
        bz.update(now)

    envelope_any = env_act_x or env_act_y

    # Matriz 2D (CMD vs ACT)
    if ticks_diff(now, t_mat) >= 0:
        t_mat = ticks_add(t_mat, MATRIX_MS)
        mx.render_xy_points(cmd_x, cmd_y, act_x, act_y, mode=mode, cas_code=cas_out["code"],
                            envelope_active=envelope_any, now_ms=now)

    # LED RGB
    if ticks_diff(now, t_rgb) >= 0:
        t_rgb = ticks_add(t_rgb, RGB_MS)
        _rgb_apply(pwm_r, pwm_g, pwm_b, mode, cas_out["code"], envelope_any, now)

    # OLED
    if ticks_diff(now, t_oled) >= 0:
        t_oled = ticks_add(t_oled, OLED_MS)
        mode_txt = mode + (" CAL" if (not js.is_calibrated()) else "")
        disp.draw_xy(
            now_ms=now,
            mode=mode_txt,
            cmd_x=cmd_x, act_x=act_x,
            cmd_y=cmd_y, act_y=act_y,
            cas_out=cas_out,
            envelope_active=envelope_any,
            alpha_used=(alpha_x if mode == MODE_NORMAL else 1.0),
        )