main.py

import time
from config import UDP_PORT, TIMEOUT_S, LOOP_HZ
from udp_rx import UdpReceiver
from led_driver import LedDriver, LED_CONNECTED, LED_LOST, LED_MANUAL, LED_CALM, LED_ALERT
from dog_driver import WaveGoDogDriver


def joystick_map(joyx: int, joyy: int):
    v = joyy / 100.0
    w = joyx / 100.0
    if abs(v) < 0.08: v = 0.0
    if abs(w) < 0.08: w = 0.0
    return v, w


def human_map(calm: float, engage: float):
    if calm > 0.7 and engage < 0.3:
        return 0.0, 0.0
    v = max(0.0, min(1.0, engage)) * 0.6
    return v, 0.0


def main():
    rx = UdpReceiver(port=UDP_PORT)
    leds = LedDriver()
    dog = WaveGoDogDriver()
    dog.connect()

    dt = 1.0 / float(LOOP_HZ)
    last_seq = None

    print(f"[OK] UDP :{UDP_PORT} timeout={TIMEOUT_S}s hz={LOOP_HZ}", flush=True)

    while True:
        pkt = rx.poll()

        if rx.seconds_since_last() > TIMEOUT_S:
            leds.set(LED_LOST)
            dog.stop()
            time.sleep(dt)
            continue

        leds.set(LED_CONNECTED)

        if pkt is None:
            time.sleep(dt)
            continue

        if pkt.seq != last_seq:
            last_seq = pkt.seq
            print(f"[PKT] seq={pkt.seq} override={pkt.override} joy=({pkt.joyx},{pkt.joyy}) calm={pkt.calm:.2f} engage={pkt.engage:.2f} emo={pkt.emotion} trig={pkt.trig}", flush=True)

        if pkt.override == 1:
            leds.set(LED_MANUAL)
            v, w = joystick_map(pkt.joyx, pkt.joyy)
            dog.set_velocity(v, w)
        else:
            if pkt.emotion.lower() in ("alert", "engaged"):
                leds.set(LED_ALERT)
            else:
                leds.set(LED_CALM)
            v, w = human_map(pkt.calm, pkt.engage)
            dog.set_velocity(v, w)

        if pkt.trig == 1:
            leds.trigger_flash()
            dog.action_jump()

        time.sleep(dt)


if __name__ == "__main__":
    main()