"""Wireless-jitter baseline: compare current-Marionette-style streaming
vs. true daemon-side playback for the same move.

Two modes, N runs each, interleaved:
  - marionette_style: Python streams one combined set_target per tick at
    50 Hz with lead-comp (mirrors what the JS app currently does on the
    data channel).
  - daemon_native: POST /move/play/recorded-move-dataset/{name}, which
    runs Backend.play_move on the robot at 100 Hz with no network in
    the inner loop. This is the target behavior of the upload-and-play
    rework -- same end state, different way of getting the move there.

For each run we capture the actual head pose / antenna positions at
50 Hz from the cached SDK state.  Output: a single summary printed to
stdout and a CSV per run under OUT_DIR for later plotting.

Usage:
    python wireless_baseline.py                 # default: head_tilt_roll
    python wireless_baseline.py <dance_name>    # any from the dataset
"""

import csv
import math
import statistics
import sys
import threading
import time
from pathlib import Path

import numpy as np
import requests

from reachy_mini import ReachyMini
from reachy_mini.motion.recorded_move import RecordedMove, RecordedMoves

DATASET = "pollen-robotics/reachy-mini-dances-library"
DANCE = sys.argv[1] if len(sys.argv) > 1 else "head_tilt_roll"
N_RUNS = 3
CAPTURE_HZ = 50.0
CAPTURE_PERIOD = 1.0 / CAPTURE_HZ
HEAD_LEAD_S = 0.205
ANTENNA_LEAD_S = 0.090
STREAM_HZ = 50.0
STREAM_PERIOD = 1.0 / STREAM_HZ

DAEMON_HOST = "http://reachy-mini.local:8000/api"

OUT_DIR = Path("/Users/remi/Downloads/wireless-baseline")
OUT_DIR.mkdir(parents=True, exist_ok=True)


def sample_move_at(move: RecordedMove, t: float):
    t_clipped = min(max(0.0, t), move.timestamps[-1] - 1e-3)
    return move.evaluate(t_clipped)


def capture_loop(mini: ReachyMini, samples: list, stop_event: threading.Event, t0: float):
    """Poll cached state at CAPTURE_HZ; append (t, head_flat16, antennas, perf_t)."""
    tick = 0
    while not stop_event.is_set():
        now = time.perf_counter()
        t = now - t0
        try:
            head = mini.get_current_head_pose()
            ant = mini.get_present_antenna_joint_positions()
            samples.append((t, head.flatten().tolist(), list(ant), now))
        except Exception as e:
            print(f"  capture read error at t={t:.3f}: {e}")
        tick += 1
        next_tick = t0 + tick * CAPTURE_PERIOD
        sleep_for = next_tick - time.perf_counter()
        if sleep_for > 0:
            time.sleep(sleep_for)


def reset_to_base(mini: ReachyMini):
    base = np.eye(4)
    mini.goto_target(base, antennas=[-0.1745, 0.1745], duration=0.8)
    time.sleep(1.0)


def head_pose_xyz_rpy(flat16: list) -> tuple:
    """Pull the translation (x,y,z) + roll/pitch/yaw from a 4×4."""
    m = np.array(flat16).reshape(4, 4)
    x, y, z = m[0, 3], m[1, 3], m[2, 3]
    # ZYX intrinsic: roll = atan2(R32,R33), pitch = asin(-R31), yaw = atan2(R21,R11)
    pitch = math.asin(-max(-1.0, min(1.0, m[2, 0])))
    if abs(math.cos(pitch)) > 1e-6:
        roll = math.atan2(m[2, 1], m[2, 2])
        yaw = math.atan2(m[1, 0], m[0, 0])
    else:
        roll = math.atan2(-m[1, 2], m[1, 1])
        yaw = 0.0
    return x, y, z, roll, pitch, yaw


def actuator_smoothness(samples: list) -> dict:
    """Pure smoothness metric -- no reference trajectory needed.

    Take per-axis position from each captured sample, compute the
    second difference (proxy for acceleration), report its std and max.
    Smooth motion has small / uniform accelerations; jerky motion
    spikes.  Independent of tracking error -- measures the actuator's
    own behavior directly.
    """
    if len(samples) < 4:
        return {"head_jerk_rms_mrad": 0.0, "ant_jerk_rms_mrad": 0.0,
                "head_jerk_max_mrad": 0.0, "ant_jerk_max_mrad": 0.0}
    head_rpy = []
    ants = []
    for _, head_flat, ant, _ in samples:
        _, _, _, r, p, y = head_pose_xyz_rpy(head_flat)
        head_rpy.append((r, p, y))
        ants.append(ant)
    def second_diff_rms(seq, axis_count):
        sq_sum = 0.0
        n = 0
        peak = 0.0
        for i in range(2, len(seq)):
            for ax in range(axis_count):
                d2 = seq[i][ax] - 2 * seq[i-1][ax] + seq[i-2][ax]
                sq_sum += d2 * d2
                if abs(d2) > peak:
                    peak = abs(d2)
                n += 1
        rms = math.sqrt(sq_sum / max(1, n))
        return rms, peak
    h_rms, h_peak = second_diff_rms(head_rpy, 3)
    a_rms, a_peak = second_diff_rms(ants, 2)
    return {
        "head_jerk_rms_mrad": h_rms * 1000,
        "head_jerk_max_mrad": h_peak * 1000,
        "ant_jerk_rms_mrad": a_rms * 1000,
        "ant_jerk_max_mrad": a_peak * 1000,
    }


def per_axis_lag_compensated_error(move: RecordedMove, samples: list) -> dict:
    """Compute commanded-vs-actual error stats with lead-comp applied.

    The marionette-style stream applies HEAD_LEAD/ANTENNA_LEAD to the
    *sampled-from-motion* time so the actuator at wall-clock t is
    expected to be at move(t).  Daemon-native doesn't lead-comp so
    actuator lags by a fixed amount.  Reporting both with and without
    lead compensation lets us compare apples to apples.
    """
    head_err = []
    ant_err = []
    head_err_compensated = []
    ant_err_compensated = []
    for t, head_flat, ant, _ in samples:
        if t < 0 or t > move.duration - 0.05:
            continue
        # uncompensated: actuator at t vs ideal at t
        head_ideal, ant_ideal, _ = sample_move_at(move, t)
        x, y, z, roll, pitch, yaw = head_pose_xyz_rpy(head_flat)
        xi, yi, zi, ri, pi_, yi_ = head_pose_xyz_rpy(head_ideal.flatten().tolist())
        head_err.append(
            math.hypot(roll - ri, math.hypot(pitch - pi_, yaw - yi_))
        )
        ant_err.append(math.hypot(ant[0] - ant_ideal[0], ant[1] - ant_ideal[1]))
        # compensated: actuator at t vs ideal at t - lead (i.e. what was commanded ~lead s ago)
        head_ideal2, _, _ = sample_move_at(move, max(0.0, t - HEAD_LEAD_S))
        _, ant_ideal2, _ = sample_move_at(move, max(0.0, t - ANTENNA_LEAD_S))
        xi, yi, zi, ri2, pi2, yi2 = head_pose_xyz_rpy(head_ideal2.flatten().tolist())
        head_err_compensated.append(
            math.hypot(roll - ri2, math.hypot(pitch - pi2, yaw - yi2))
        )
        ant_err_compensated.append(
            math.hypot(ant[0] - ant_ideal2[0], ant[1] - ant_ideal2[1])
        )
    def stats(xs):
        if not xs:
            return (0.0, 0.0, 0.0)
        return (statistics.mean(xs), statistics.median(xs), max(xs))
    return {
        "head_uncomp": stats(head_err),
        "ant_uncomp": stats(ant_err),
        "head_lag_comp": stats(head_err_compensated),
        "ant_lag_comp": stats(ant_err_compensated),
        "n": len(head_err),
    }


def stream_period_stats(send_times_perf: list) -> tuple:
    """Tick interval mean/std/max from the streamer's perf_counter times."""
    if len(send_times_perf) < 2:
        return (0.0, 0.0, 0.0)
    diffs = [send_times_perf[i+1] - send_times_perf[i] for i in range(len(send_times_perf)-1)]
    return (statistics.mean(diffs), statistics.stdev(diffs) if len(diffs) > 1 else 0.0, max(diffs))


def write_csv(path: Path, move: RecordedMove, samples: list):
    headers = ["t_s", "act_x", "act_y", "act_z", "act_roll", "act_pitch", "act_yaw", "act_ant_r", "act_ant_l",
               "cmd_x", "cmd_y", "cmd_z", "cmd_roll", "cmd_pitch", "cmd_yaw", "cmd_ant_r", "cmd_ant_l"]
    with open(path, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(headers)
        for t, head_flat, ant, _ in samples:
            if t < 0 or t > move.duration + 0.2:
                continue
            ax, ay, az, ar, ap, aw = head_pose_xyz_rpy(head_flat)
            head_cmd, ant_cmd, _ = sample_move_at(move, t)
            cx, cy, cz, cr, cp, cw = head_pose_xyz_rpy(head_cmd.flatten().tolist())
            w.writerow([f"{t:.4f}",
                        f"{ax:.5f}", f"{ay:.5f}", f"{az:.5f}", f"{ar:.5f}", f"{ap:.5f}", f"{aw:.5f}",
                        f"{ant[0]:.5f}", f"{ant[1]:.5f}",
                        f"{cx:.5f}", f"{cy:.5f}", f"{cz:.5f}", f"{cr:.5f}", f"{cp:.5f}", f"{cw:.5f}",
                        f"{ant_cmd[0]:.5f}", f"{ant_cmd[1]:.5f}"])


def run_marionette_style(mini: ReachyMini, move: RecordedMove, run_idx: int) -> dict:
    samples = []
    send_times = []
    stop = threading.Event()
    t0 = time.perf_counter()
    capture_thread = threading.Thread(
        target=capture_loop, args=(mini, samples, stop, t0), daemon=True,
    )
    capture_thread.start()
    try:
        tick = 0
        while True:
            now = time.perf_counter()
            t = now - t0
            if t > move.duration:
                break
            t_head = min(t + HEAD_LEAD_S, move.duration - 1e-3)
            t_ant = min(t + ANTENNA_LEAD_S, move.duration - 1e-3)
            head_for_head, _, body_yaw_h = sample_move_at(move, t_head)
            _, ant_for_ant, _ = sample_move_at(move, t_ant)
            mini.set_target(
                head=head_for_head,
                antennas=list(ant_for_ant),
                body_yaw=float(body_yaw_h),
            )
            send_times.append(now)
            tick += 1
            next_tick = t0 + tick * STREAM_PERIOD
            sleep_for = next_tick - time.perf_counter()
            if sleep_for > 0:
                time.sleep(sleep_for)
        time.sleep(0.3)
    finally:
        stop.set()
        capture_thread.join(timeout=1.0)
    out = OUT_DIR / f"{DANCE}-marionette_style-run{run_idx:02d}.csv"
    write_csv(out, move, samples)
    return {
        "mode": "marionette_style",
        "csv": out,
        "stream_stats": stream_period_stats(send_times),
        "err": per_axis_lag_compensated_error(move, samples),
        "smooth": actuator_smoothness(samples),
        "frames_sent": len(send_times),
        "frames_captured": len(samples),
    }


def run_daemon_native(mini: ReachyMini, move: RecordedMove, dance_name: str, run_idx: int) -> dict:
    samples = []
    stop = threading.Event()
    t0 = time.perf_counter()
    capture_thread = threading.Thread(
        target=capture_loop, args=(mini, samples, stop, t0), daemon=True,
    )
    capture_thread.start()
    try:
        # Kick off the daemon-side playback via REST. It runs
        # Backend.play_move on the robot at 100 Hz with no network in
        # the inner loop.
        r = requests.post(
            f"{DAEMON_HOST}/move/play/recorded-move-dataset/{DATASET}/{dance_name}",
            timeout=5,
        )
        r.raise_for_status()
        move_uuid = r.json()["uuid"]
        # Wait for the daemon to finish (poll the running-moves list).
        deadline = time.perf_counter() + move.duration + 5.0
        while time.perf_counter() < deadline:
            time.sleep(0.05)
            try:
                running = requests.get(f"{DAEMON_HOST}/move/running", timeout=2).json()
                if not any(m["uuid"] == move_uuid for m in running):
                    break
            except Exception:
                pass
        time.sleep(0.3)
    finally:
        stop.set()
        capture_thread.join(timeout=1.0)
    out = OUT_DIR / f"{DANCE}-daemon_native-run{run_idx:02d}.csv"
    write_csv(out, move, samples)
    return {
        "mode": "daemon_native",
        "csv": out,
        "stream_stats": (0.01, 0.0, 0.0),  # not applicable
        "err": per_axis_lag_compensated_error(move, samples),
        "smooth": actuator_smoothness(samples),
        "frames_sent": int(move.duration * 100),  # daemon ticks at 100 Hz
        "frames_captured": len(samples),
    }


def fmt_stats(label: str, stats: tuple) -> str:
    mean, median, mx = stats
    return f"{label}: mean={mean*1000:.1f} ms, median={median*1000:.1f} ms, max={mx*1000:.1f} ms"


def summarize(results: list):
    print("\n" + "=" * 70)
    print(f"SUMMARY -- dance={DANCE}, {N_RUNS} runs per mode")
    print("=" * 70)
    by_mode: dict[str, list] = {"marionette_style": [], "daemon_native": []}
    for r in results:
        by_mode[r["mode"]].append(r)
    for mode, rs in by_mode.items():
        if not rs:
            continue
        # Aggregate err stats across runs.
        def agg(key, idx):
            vals = [r["err"][key][idx] for r in rs]
            return statistics.mean(vals) if vals else 0.0
        # Stream interval stats (mean of means across runs)
        if mode == "marionette_style":
            tick_means = [r["stream_stats"][0] for r in rs]
            tick_stds = [r["stream_stats"][1] for r in rs]
            tick_maxes = [r["stream_stats"][2] for r in rs]
            tick_summary = (
                f"  send-tick: mean={statistics.mean(tick_means)*1000:.1f} ms, "
                f"std={statistics.mean(tick_stds)*1000:.1f} ms, "
                f"max={max(tick_maxes)*1000:.1f} ms (target 20 ms)"
            )
        else:
            tick_summary = "  send-tick: n/a (daemon-side loop)"

        print(f"\n[{mode}]  runs={len(rs)}")
        print(tick_summary)
        # Smoothness (no reference trajectory needed -- pure actuator
        # behavior).  Lower is smoother.
        sh = statistics.mean([r["smooth"]["head_jerk_rms_mrad"] for r in rs])
        sh_max = max([r["smooth"]["head_jerk_max_mrad"] for r in rs])
        sa = statistics.mean([r["smooth"]["ant_jerk_rms_mrad"] for r in rs])
        sa_max = max([r["smooth"]["ant_jerk_max_mrad"] for r in rs])
        print(f"  smoothness (lower is better):")
        print(f"    head rpy jerk rms={sh:.2f} mrad,  max={sh_max:.2f} mrad")
        print(f"    antennas jerk rms={sa:.2f} mrad,  max={sa_max:.2f} mrad")
        print(f"  tracking error (sanity, not the primary metric):")
        print(f"    lag-comp head={agg('head_lag_comp', 0)*1000:.1f} mrad / "
              f"ant={agg('ant_lag_comp', 0)*1000:.1f} mrad")
    print()


def main():
    print(f"Loading dance: {DANCE} from {DATASET}")
    library = RecordedMoves(DATASET)
    if DANCE not in library.list_moves():
        print(f"Dance '{DANCE}' not found. Available: {library.list_moves()}")
        sys.exit(2)
    move = library.get(DANCE)
    print(f"  duration: {move.duration:.2f} s, frames: {len(move.timestamps)}")

    print("Connecting to robot…")
    with ReachyMini() as mini:
        print("Connected. Resetting to base…")
        reset_to_base(mini)

        results = []
        for run_idx in range(1, N_RUNS + 1):
            print(f"\n--- Run {run_idx}/{N_RUNS}: marionette_style ---")
            r = run_marionette_style(mini, move, run_idx)
            print(f"  sent {r['frames_sent']} frames, captured {r['frames_captured']}")
            results.append(r)
            reset_to_base(mini)

            print(f"--- Run {run_idx}/{N_RUNS}: daemon_native ---")
            r = run_daemon_native(mini, move, DANCE, run_idx)
            print(f"  captured {r['frames_captured']} frames")
            results.append(r)
            reset_to_base(mini)

        summarize(results)

    print(f"\nCSVs in {OUT_DIR}")


if __name__ == "__main__":
    main()