#!/usr/bin/env python3
"""Measure whether cancelling a gesture mid-play interpolates or SNAPS, and write
PREEMPTION.md.

A child interrupts constantly, so preemption must ramp smoothly to idle, not jump.
This probe (run ON THE ROBOT, or in sim with --sim) does two runs:

  1. RAW: start a move, call `cancel_move()` mid-play, and sample joint positions
     across the cancel. A large single-sample jump == the SDK snaps.
  2. WRAPPED: same, but let IdleMotion's ramp-to-idle take over after the cancel.

It reports the worst per-sample joint delta for each and writes the verdict into
PREEMPTION.md. Bounded, smooth deltas == good; a spike == the ramp-down wrapper is
doing real work.
"""

from __future__ import annotations

import argparse
import statistics
import threading
import time
from pathlib import Path

import numpy as np

RESULTS_START = "<!-- RESULTS:START -->"
RESULTS_END = "<!-- RESULTS:END -->"
PREEMPTION_MD = Path(__file__).resolve().parent.parent / "PREEMPTION.md"

SAMPLE_HZ = 100.0


def sample_joints(mini, out: list[np.ndarray], stop: threading.Event) -> None:
    get = getattr(mini, "get_current_joint_positions", None)
    if get is None:
        return
    dt = 1.0 / SAMPLE_HZ
    while not stop.is_set():
        try:
            out.append(np.asarray(get(), dtype=float))
        except Exception:  # noqa: BLE001
            pass
        time.sleep(dt)


def max_step(samples: list[np.ndarray]) -> float:
    """Largest single-sample L-inf joint delta (rad) across the capture."""
    if len(samples) < 2:
        return float("nan")
    steps = [np.max(np.abs(b - a)) for a, b in zip(samples, samples[1:]) if a.shape == b.shape]
    return float(max(steps)) if steps else float("nan")


def run_once(mini, idle, use_wrapper: bool) -> float:
    from hermes_voice.gestures import GestureLibrary

    lib = GestureLibrary(mini, idle)
    samples: list[np.ndarray] = []
    stop = threading.Event()
    sampler = threading.Thread(target=sample_joints, args=(mini, samples, stop), daemon=True)
    sampler.start()
    lib.request("celebrate")           # a longer, expressive move
    time.sleep(0.6)                    # let it get going
    if use_wrapper:
        lib.request("nod_gentle")      # supersede -> cancel + idle ramp-to-idle
    else:
        cancel = getattr(mini, "cancel_move", None)
        if cancel:
            cancel()
    time.sleep(1.2)
    stop.set()
    sampler.join(timeout=1.0)
    lib.stop()
    return max_step(samples)


def write_results(raw: float, wrapped: float) -> None:
    verdict = "inconclusive"
    if not np.isnan(raw) and not np.isnan(wrapped):
        verdict = "wrapper smooths the cancel" if wrapped < raw * 0.75 else \
                  "SDK cancel already smooth" if raw < 0.05 else "check tuning"
    block = (
        f"{RESULTS_START}\n"
        f"| run | worst single-sample joint step (rad) |\n"
        f"|-----|--------------------------------------|\n"
        f"| raw `cancel_move()` | {raw:.4f} |\n"
        f"| wrapped (ramp-to-idle) | {wrapped:.4f} |\n\n"
        f"**Verdict:** {verdict}. (Lower is smoother; a large raw step means the SDK "
        f"snaps and the ramp-down wrapper is required.)\n"
        f"{RESULTS_END}"
    )
    if PREEMPTION_MD.exists() and RESULTS_START in PREEMPTION_MD.read_text():
        text = PREEMPTION_MD.read_text()
        PREEMPTION_MD.write_text(text.split(RESULTS_START)[0] + block + text.split(RESULTS_END)[1])
        print(f"Wrote results into {PREEMPTION_MD.name}")
    print(f"raw={raw:.4f} rad  wrapped={wrapped:.4f} rad  ->  {verdict}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--sim", action="store_true")
    args = ap.parse_args()
    from reachy_mini import ReachyMini

    from hermes_voice.idle import IdleMotion

    # "no_media" = motion only: skip camera/audio/WebRTC init entirely, so the
    # probe connects fast and never hangs on media bring-up (it only needs joints).
    backend = "default" if args.sim else "no_media"
    with ReachyMini(media_backend=backend, use_sim=args.sim) as mini:
        try:
            mini.enable_motors()
        except Exception:  # noqa: BLE001
            pass
        idle = IdleMotion(mini)
        idle.start()
        try:
            raw = run_once(mini, idle, use_wrapper=False)
            time.sleep(1.0)
            wrapped = run_once(mini, idle, use_wrapper=True)
        finally:
            idle.stop()
    write_results(raw, wrapped)


if __name__ == "__main__":
    main()