#!/usr/bin/env python3
"""Test antenna collisions on Reachy Mini.

Standalone script (no Marionette). Commands the left antenna to collide
with the right antenna at known timestamps, using non-periodic intervals
so detection alignment is unambiguous.

Collision setup:
  - Right antenna fixed at -0.68 rad
  - Left antenna moves from 0.0 to 0.70 rad (collision at ~0.60 rad)
  - Hold 100ms, then return to rest
  - Low PID + flexible antennas = safe, produces audible click

Records present antenna positions at 50Hz on the robot side (no network
latency), saves to JSON, SCPs back, and analyzes to detect collision times
from the trajectory.

Usage:
    python tests/test_antenna_collision.py [--host reachy-mini.local]
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
import tempfile
import time
from pathlib import Path

import numpy as np

# Collision parameters
RIGHT_REST = -0.68    # right antenna fixed position (rad)
LEFT_REST = 0.0       # left antenna rest position (rad)
LEFT_COLLISION = 0.70 # left antenna collision target (past contact at ~0.60)
HOLD_DURATION = 0.2   # seconds to hold at collision position

# Non-periodic collision times — gaps are all different (1.3, 1.7, 2.3, 3.1s)
# so there's no ambiguity when matching detected events to expected events.
COLLISION_TIMES = [1.0, 2.3, 4.0, 6.3, 9.4]

ROBOT_USER = "pollen"
ROBOT_PYTHON = "/venvs/apps_venv/bin/python"
REMOTE_RESULTS = "/tmp/collision_positions.json"

ROBOT_COLLISION_SCRIPT = """\
import numpy as np
import os, sys, time, json

collision_times = json.loads(sys.argv[1])
right_rest = float(sys.argv[2])
left_rest = float(sys.argv[3])
left_collision = float(sys.argv[4])
hold_duration = float(sys.argv[5])
results_path = sys.argv[6]

print(f"robot: connecting to ReachyMini", flush=True)
from reachy_mini import ReachyMini
from reachy_mini.utils import create_head_pose
r = ReachyMini(media_backend="no_media")

# Go to rest position
r.goto_target(create_head_pose(), antennas=[left_rest, right_rest], duration=1.0)
time.sleep(1.5)

print(f"robot: rest position — left={left_rest}, right={right_rest}", flush=True)
print(f"robot: will collide at times: {collision_times}", flush=True)
print(f"robot: collision target: left={left_collision}, hold={hold_duration}s", flush=True)

# State machine
DT = 0.02  # 50Hz update rate
total_duration = max(collision_times) + hold_duration + 1.0
n_steps = int(total_duration / DT)

# Build timeline: for each step, determine left antenna target
left_targets = np.full(n_steps, left_rest, dtype=np.float64)
for ct in collision_times:
    start_step = int(ct / DT)
    end_step = int((ct + hold_duration) / DT)
    end_step = min(end_step, n_steps)
    left_targets[start_step:end_step] = left_collision

# Recording arrays
timestamps = []
left_present = []
right_present = []
left_target_log = []

print("robot: MARK_START", flush=True)
t0 = time.monotonic()

for i in range(n_steps):
    left = float(left_targets[i])
    r.set_target(
        head=np.eye(4),
        body_yaw=0.0,
        antennas=np.array([left, right_rest]),
    )

    # Read present position
    pos = r.get_present_antenna_joint_positions()
    elapsed = time.monotonic() - t0
    timestamps.append(elapsed)
    left_present.append(pos[0])
    right_present.append(pos[1])
    left_target_log.append(left)

    # Real-time pacing
    target_time = (i + 1) * DT
    now = time.monotonic() - t0
    if target_time > now:
        time.sleep(target_time - now)

elapsed = time.monotonic() - t0
print(f"robot: finished {n_steps} steps in {elapsed:.3f}s (expected {total_duration:.3f}s)", flush=True)

# Save results
results = {
    "collision_times": collision_times,
    "left_collision_target": left_collision,
    "right_rest": right_rest,
    "hold_duration": hold_duration,
    "timestamps": timestamps,
    "left_present": left_present,
    "right_present": right_present,
    "left_target": left_target_log,
}
with open(results_path, "w") as f:
    json.dump(results, f)
print(f"robot: saved {len(timestamps)} samples to {results_path}", flush=True)

# Return to rest
r.goto_target(create_head_pose(), antennas=[left_rest, right_rest], duration=1.0)
time.sleep(1.5)
print("robot: done", flush=True)
os._exit(0)
"""


def scp_to_robot(local_path: Path, remote_path: str, host: str) -> None:
    """Copy a file to the robot via SCP."""
    target = f"{ROBOT_USER}@{host}:{remote_path}"
    result = subprocess.run(
        ["scp", "-o", "ConnectTimeout=5", str(local_path), target],
        capture_output=True, text=True, timeout=15,
    )
    if result.returncode != 0:
        raise RuntimeError(f"SCP failed: {result.stderr}")
    print(f"  Copied to {target}")


def scp_from_robot(remote_path: str, local_path: Path, host: str) -> None:
    """Copy a file from the robot via SCP."""
    source = f"{ROBOT_USER}@{host}:{remote_path}"
    result = subprocess.run(
        ["scp", "-o", "ConnectTimeout=5", source, str(local_path)],
        capture_output=True, text=True, timeout=15,
    )
    if result.returncode != 0:
        raise RuntimeError(f"SCP failed: {result.stderr}")
    print(f"  Copied from {source}")


def run_on_robot(host: str) -> subprocess.Popen:
    """SCP script to robot and start collision sequence (non-blocking)."""
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
        f.write(ROBOT_COLLISION_SCRIPT)
        local_script = Path(f.name)

    remote_script = "/tmp/collision_test.py"
    try:
        scp_to_robot(local_script, remote_script, host)
    finally:
        local_script.unlink()

    args_str = (
        f"{ROBOT_PYTHON} {remote_script} "
        f"'{json.dumps(COLLISION_TIMES)}' "
        f"{RIGHT_REST} {LEFT_REST} {LEFT_COLLISION} {HOLD_DURATION} "
        f"{REMOTE_RESULTS}"
    )
    proc = subprocess.Popen(
        ["ssh", "-o", "ConnectTimeout=5", f"{ROBOT_USER}@{host}", args_str],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
    )
    return proc


def detect_collisions(data: dict) -> list[dict]:
    """Detect collision events from present position trajectory.

    For each commanded collision, finds the position peak time (when the
    antenna reached its maximum excursion and reversed). This is the moment
    of collision impact or max overshoot.

    Returns list of dicts with detection details per event.
    """
    ts = np.array(data["timestamps"])
    left_pos = np.array(data["left_present"])
    left_tgt = np.array(data["left_target"])

    # Find rising edges in target (rest → collision command)
    tgt_diff = np.diff(left_tgt)
    command_indices = np.where(tgt_diff > 0.3)[0]

    events = []
    for cmd_idx in command_indices:
        cmd_time = ts[cmd_idx]

        # Look 0-500ms after command for position peak
        window_start = cmd_idx
        window_end = min(cmd_idx + 25, len(ts))  # 500ms at 50Hz
        if window_end <= window_start + 2:
            continue

        window_pos = left_pos[window_start:window_end]
        window_ts = ts[window_start:window_end]

        # Position peak = moment of max excursion (collision or overshoot)
        peak_idx = np.argmax(window_pos)
        peak_time = window_ts[peak_idx]
        peak_pos = window_pos[peak_idx]

        # Check for velocity stall (plateau = physical stop at collision)
        dt = np.diff(window_ts)
        dt[dt < 1e-6] = 1e-6
        vel = np.diff(window_pos) / dt

        # A stall is consecutive near-zero velocity while position > 0.1
        stall_time = None
        for j in range(len(vel) - 1):
            if (abs(vel[j]) < 1.0 and abs(vel[j + 1]) < 1.0
                    and window_pos[j + 1] > 0.1):
                stall_time = window_ts[j + 1]
                break

        events.append({
            "cmd_time": cmd_time,
            "peak_time": peak_time,
            "peak_pos": peak_pos,
            "latency_ms": (peak_time - cmd_time) * 1000,
            "stall_time": stall_time,
        })

    return events


def plot_results(data: dict, events: list[dict], output_path: Path) -> None:
    """Generate a plot of antenna trajectories with collision markers."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    ts = np.array(data["timestamps"])
    left_pos = np.array(data["left_present"])
    right_pos = np.array(data["right_present"])
    left_tgt = np.array(data["left_target"])
    expected_times = data["collision_times"]

    fig, ax = plt.subplots(figsize=(14, 6))

    # Plot antenna trajectories
    ax.plot(ts, left_pos, "b-", linewidth=1.5, label="Left antenna (present)", zorder=3)
    ax.plot(ts, right_pos, "r-", linewidth=1.5, label="Right antenna (present)", zorder=3)
    ax.plot(ts, left_tgt, "b--", linewidth=0.8, alpha=0.4, label="Left antenna (target)")

    # Expected collision times (green dashed)
    for i, ct in enumerate(expected_times):
        label = "Expected collision" if i == 0 else None
        ax.axvline(ct, color="green", linestyle="--", linewidth=1.5, alpha=0.7, label=label)

    # Detected collision times (red solid)
    for i, ev in enumerate(events):
        label = "Detected (peak)" if i == 0 else None
        ax.axvline(ev["peak_time"], color="red", linestyle="-", linewidth=1.5, alpha=0.7, label=label)
        # Annotate with peak position
        ax.annotate(
            f'{ev["peak_pos"]:.2f}r\n{ev["latency_ms"]:.0f}ms',
            xy=(ev["peak_time"], ev["peak_pos"]),
            xytext=(10, 10), textcoords="offset points",
            fontsize=8, color="red",
            arrowprops=dict(arrowstyle="->", color="red", lw=0.8),
        )

    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Position (rad)")
    ax.set_title("Antenna Collision Test — Present Positions + Detection")
    ax.legend(loc="upper right", fontsize=9)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(ts[0] - 0.2, ts[-1] + 0.2)

    fig.tight_layout()
    fig.savefig(str(output_path), dpi=150)
    plt.close(fig)
    print(f"  Plot saved to {output_path}")


def analyze_results(data: dict) -> dict:
    """Analyze collision position data and print report."""
    ts = np.array(data["timestamps"])
    left_pos = np.array(data["left_present"])
    left_tgt = np.array(data["left_target"])
    expected_times = data["collision_times"]

    print(f"\n{'='*60}")
    print("Collision Detection Analysis")
    print(f"{'='*60}")
    print(f"  Samples: {len(ts)} at ~{1/np.mean(np.diff(ts)):.0f}Hz")
    print(f"  Duration: {ts[-1]:.2f}s")
    print(f"  Left antenna range: [{left_pos.min():.3f}, {left_pos.max():.3f}] rad")
    print(f"  Expected collisions: {len(expected_times)} at {expected_times}")

    # Detect collisions
    events = detect_collisions(data)

    # Per-collision detail
    print(f"\n--- Per-collision detail ---")
    for i, ev in enumerate(events):
        stall_info = f", stall at t={ev['stall_time']:.3f}s" if ev["stall_time"] else ""
        print(f"  Collision {i+1}: cmd t={ev['cmd_time']:.1f}s → "
              f"peak {ev['peak_pos']:.3f} rad at t={ev['peak_time']:.3f}s "
              f"(latency {ev['latency_ms']:.0f}ms{stall_info})")

    # Match detected peaks to expected times
    detected_peak_times = [ev["peak_time"] for ev in events]
    print(f"\n--- Matching ---")
    matched = 0
    for i, et in enumerate(expected_times):
        if not detected_peak_times:
            print(f"  Expected {et:.1f}s: NO MATCH")
            continue
        nearest = min(detected_peak_times, key=lambda d: abs(d - et))
        offset_ms = (nearest - et) * 1000
        ok = abs(nearest - et) < 0.3
        if ok:
            matched += 1
        status = "OK" if ok else "MISS"
        print(f"  Expected {et:.1f}s → peak at {nearest:.3f}s "
              f"(offset {offset_ms:+.0f}ms) [{status}]")

    success = matched >= len(expected_times) - 1
    print(f"\n  {'PASS' if success else 'FAIL'}: "
          f"{matched}/{len(expected_times)} collisions matched")

    # Generate plot
    plot_path = Path("tests/collision_trajectory.png")
    plot_results(data, events, plot_path)

    return {
        "success": success,
        "expected": len(expected_times),
        "detected": len(events),
        "matched": matched,
        "events": events,
    }


def main():
    parser = argparse.ArgumentParser(description="Antenna collision test")
    parser.add_argument("--host", default="reachy-mini.local", help="Robot hostname/IP")
    args = parser.parse_args()

    print(f"\n{'='*60}")
    print("Antenna Collision Test")
    print(f"{'='*60}")
    print(f"  Right antenna fixed at {RIGHT_REST} rad")
    print(f"  Left antenna: rest={LEFT_REST}, collision={LEFT_COLLISION} rad")
    print(f"  Collision times: {COLLISION_TIMES}")
    gaps = [COLLISION_TIMES[i+1] - COLLISION_TIMES[i] for i in range(len(COLLISION_TIMES)-1)]
    print(f"  Gaps between collisions: {[f'{g:.1f}s' for g in gaps]}")
    print(f"  Hold duration: {HOLD_DURATION}s\n")

    # Stop any running app
    print("[1/4] Stopping any running app on robot...")
    subprocess.run(
        ["ssh", "-o", "ConnectTimeout=5", f"{ROBOT_USER}@{args.host}",
         "curl -sf -X POST http://127.0.0.1:8000/api/apps/stop-current-app >/dev/null 2>&1 || true"],
        capture_output=True, timeout=10,
    )
    time.sleep(1)
    print("  Done")

    # Run collisions (with position recording)
    print("\n[2/4] Running collision sequence + recording positions...")
    proc = run_on_robot(args.host)

    # Stream robot output in real time
    print("\n--- Robot output ---")
    for line in iter(proc.stdout.readline, ""):
        line = line.rstrip()
        if line:
            print(f"  {line}")
    proc.wait()
    print("--- End robot output ---")

    if proc.returncode != 0:
        print(f"\nFAILED — Return code: {proc.returncode}")
        return 1

    # SCP results back
    print("\n[3/4] Fetching position data from robot...")
    local_results = Path("tests/collision_positions.json")
    scp_from_robot(REMOTE_RESULTS, local_results, args.host)
    print(f"  Saved to {local_results}")

    # Load and analyze
    print("\n[4/4] Analyzing collision trajectory...")
    with open(local_results) as f:
        data = json.load(f)

    result = analyze_results(data)

    print(f"\n{'='*60}")
    if result["success"]:
        print("RESULT: PASS — Collisions detected from present position trajectory")
    else:
        print("RESULT: FAIL — Could not reliably detect collisions")
    print(f"{'='*60}\n")

    return 0 if result["success"] else 1


if __name__ == "__main__":
    sys.exit(main())