tests/test_beep_collision_sync.py

#!/usr/bin/env python3
"""Combined beep + collision sync test.

The robot plays beeps through its speaker AND performs antenna collisions,
with exactly 1.0s between each beep and its corresponding collision.
The laptop mic records everything. Since both events are detected from
the same mic recording, the measured interval is free of cross-clock bias.

If audio and motion are perfectly synced, each beep-collision pair should
be exactly 1.0s apart in the mic recording. Deviations measure the true
audio-motion sync error.

Beep times (non-periodic, gaps 1.3/1.7/2.3/3.1s):
    [1.0, 2.3, 4.0, 6.3, 9.4]
Collision times (each beep + 1.0s):
    [2.0, 3.3, 5.0, 7.3, 10.4]

Usage:
    python tests/test_beep_collision_sync.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
import sounddevice as sd
import soundfile as sf

sys.path.insert(0, str(Path(__file__).parent))
from audio_analysis import detect_beep_onsets, detect_transient_onsets

# Timing — non-periodic gaps (1.3, 1.7, 2.3, 3.1s)
BEEP_TIMES = [1.0, 2.3, 4.0, 6.3, 9.4]
BEEP_COLLISION_OFFSET = 1.0  # seconds between beep and its collision
COLLISION_TIMES = [t + BEEP_COLLISION_OFFSET for t in BEEP_TIMES]

# Audio parameters
BEEP_FREQ = 2000.0
BEEP_DURATION = 0.2
BEEP_AMPLITUDE = 0.9
ROBOT_SR = 16000

# Collision parameters
RIGHT_REST = -0.68
LEFT_REST = 0.0
LEFT_COLLISION = 0.70
HOLD_DURATION = 0.2

ROBOT_USER = "pollen"
ROBOT_PYTHON = "/venvs/apps_venv/bin/python"
LAPTOP_SR = 48000
MIC_DURATION = 30.0
REMOTE_RESULTS = "/tmp/beep_collision_results.json"

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

beep_times = json.loads(sys.argv[1])
collision_times = json.loads(sys.argv[2])
right_rest = float(sys.argv[3])
left_rest = float(sys.argv[4])
left_collision = float(sys.argv[5])
hold_duration = float(sys.argv[6])
wav_path = sys.argv[7]
results_path = sys.argv[8]
beep_freq = float(sys.argv[9])
beep_duration = float(sys.argv[10])
beep_amplitude = float(sys.argv[11])

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

robot_sr = int(r.media.get_output_audio_samplerate() or 16000)
print(f"robot: audio sr={robot_sr}", flush=True)

# Generate beep audio
total_duration = max(collision_times) + hold_duration + 1.0
n_total = int(robot_sr * total_duration)
audio = np.zeros(n_total, dtype=np.float32)
for bt in beep_times:
    start = int(bt * robot_sr)
    n_beep = int(beep_duration * robot_sr)
    if start + n_beep > n_total:
        continue
    t_arr = np.arange(n_beep, dtype=np.float32) / robot_sr
    beep = beep_amplitude * np.sin(2 * np.pi * beep_freq * t_arr).astype(np.float32)
    fade = int(0.005 * robot_sr)
    if fade > 0 and 2 * fade < n_beep:
        beep[:fade] *= np.linspace(0, 1, fade, dtype=np.float32)
        beep[-fade:] *= np.linspace(1, 0, fade, dtype=np.float32)
    audio[start:start + n_beep] += beep

print(f"robot: generated {total_duration:.1f}s audio with {len(beep_times)} beeps", flush=True)

# Build collision timeline at 50Hz
DT = 0.02
n_steps = int(total_duration / DT)
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

# Go to rest
r.goto_target(create_head_pose(), antennas=[left_rest, right_rest], duration=1.0)
time.sleep(1.5)
print(f"robot: beeps at {beep_times}", flush=True)
print(f"robot: collisions at {collision_times}", flush=True)

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

# Start audio playback
r.media.start_playing()
r.media.push_audio_sample(np.zeros(160, dtype=np.float32))
time.sleep(0.05)

# Audio chunk tracking
chunk_size = int(robot_sr * DT)  # 20ms audio chunks match motion DT
audio_idx = 0

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

for i in range(n_steps):
    # Push audio chunk
    chunk_start = i * chunk_size
    chunk_end = chunk_start + chunk_size
    if chunk_end <= len(audio):
        r.media.push_audio_sample(audio[chunk_start:chunk_end])

    # Set antenna target
    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", flush=True)

# Drain audio buffer and stop
time.sleep(0.5)
r.media.stop_playing()

# Save results
results = {
    "beep_times": beep_times,
    "collision_times": collision_times,
    "beep_collision_offset": collision_times[0] - beep_times[0],
    "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)

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:
    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:
    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 start_robot(host: str) -> subprocess.Popen:
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
        f.write(ROBOT_SCRIPT)
        local_script = Path(f.name)

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

    args_str = (
        f"{ROBOT_PYTHON} {remote_script} "
        f"'{json.dumps(BEEP_TIMES)}' "
        f"'{json.dumps(COLLISION_TIMES)}' "
        f"{RIGHT_REST} {LEFT_REST} {LEFT_COLLISION} {HOLD_DURATION} "
        f"/dev/null "  # wav_path unused, generated inline
        f"{REMOTE_RESULTS} "
        f"{BEEP_FREQ} {BEEP_DURATION} {BEEP_AMPLITUDE}"
    )
    proc = subprocess.Popen(
        ["ssh", "-o", "ConnectTimeout=5", f"{ROBOT_USER}@{host}", args_str],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
    )
    return proc


def plot_combined(
    mic_audio: np.ndarray,
    mic_sr: int,
    mic_start: float,
    mark_start: float,
    robot_data: dict,
    detected_beeps: list[float],
    detected_collisions: list[float],
    pairs: list[dict],
    output_path: Path,
) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    mic_t = np.arange(len(mic_audio)) / mic_sr
    robot_offset = mark_start - mic_start

    robot_ts = np.array(robot_data["timestamps"])
    left_pos = np.array(robot_data["left_present"])
    right_pos = np.array(robot_data["right_present"])
    left_tgt = np.array(robot_data["left_target"])

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(18, 10), sharex=True)

    # --- Top: Mic waveform ---
    ax1.plot(mic_t, mic_audio, "k-", linewidth=0.3, alpha=0.5)
    ax1.set_ylabel("Mic amplitude")
    ax1.set_title("Beep + Collision Sync Test — Laptop Mic Recording")
    ax1.grid(True, alpha=0.3)

    # Detected beeps (blue)
    for i, bt in enumerate(detected_beeps):
        label = "Detected beep" if i == 0 else None
        ax1.axvline(bt, color="blue", linestyle="-", linewidth=1.2, alpha=0.7, label=label)

    # Detected collisions (red)
    for i, ct in enumerate(detected_collisions):
        label = "Detected collision" if i == 0 else None
        ax1.axvline(ct, color="red", linestyle="-", linewidth=1.2, alpha=0.7, label=label)

    # Annotate pairs
    for p in pairs:
        mid = (p["beep_mic_t"] + p["collision_mic_t"]) / 2
        ax1.annotate(
            f'{p["interval_ms"]:.0f}ms',
            xy=(mid, ax1.get_ylim()[1] * 0.8),
            ha="center", fontsize=9, color="purple", fontweight="bold",
            bbox=dict(boxstyle="round,pad=0.2", facecolor="lightyellow", alpha=0.8),
        )

    ax1.legend(loc="upper right", fontsize=9)

    # --- Bottom: Robot trajectory ---
    ax2.plot(robot_ts + robot_offset, left_pos, "b-", linewidth=1.5, label="Left antenna (present)")
    ax2.plot(robot_ts + robot_offset, right_pos, "r-", linewidth=1.5, label="Right antenna (present)")
    ax2.plot(robot_ts + robot_offset, left_tgt, "b--", linewidth=0.8, alpha=0.4, label="Left antenna (target)")

    # Expected command times (robot clock → mic clock)
    for i, bt in enumerate(BEEP_TIMES):
        mic_bt = robot_offset + bt
        label = "Beep cmd" if i == 0 else None
        ax2.axvline(mic_bt, color="blue", linestyle="--", linewidth=1.0, alpha=0.5, label=label)
    for i, ct in enumerate(COLLISION_TIMES):
        mic_ct = robot_offset + ct
        label = "Collision cmd" if i == 0 else None
        ax2.axvline(mic_ct, color="red", linestyle="--", linewidth=1.0, alpha=0.5, label=label)

    # Detected events on trajectory panel too
    for bt in detected_beeps:
        ax2.axvline(bt, color="blue", linestyle="-", linewidth=0.8, alpha=0.4)
    for ct in detected_collisions:
        ax2.axvline(ct, color="red", linestyle="-", linewidth=0.8, alpha=0.4)

    ax2.set_xlabel("Time since mic start (s)")
    ax2.set_ylabel("Position (rad)")
    ax2.set_title("Robot Antenna Trajectory (aligned to mic clock)")
    ax2.legend(loc="upper right", fontsize=9)
    ax2.grid(True, alpha=0.3)

    # Zoom to active region
    active_start = robot_offset - 0.5
    active_end = robot_offset + max(COLLISION_TIMES) + 2.0
    ax1.set_xlim(active_start, active_end)

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


def main():
    parser = argparse.ArgumentParser(description="Beep + collision sync test")
    parser.add_argument("--host", default="reachy-mini.local")
    args = parser.parse_args()

    print(f"\n{'='*60}")
    print("Beep + Collision Sync Test")
    print(f"{'='*60}")
    print(f"  Beep times:      {BEEP_TIMES}")
    print(f"  Collision times:  {COLLISION_TIMES}")
    print(f"  Expected interval: {BEEP_COLLISION_OFFSET:.1f}s (beep → collision)")
    gaps = [BEEP_TIMES[i+1] - BEEP_TIMES[i] for i in range(len(BEEP_TIMES)-1)]
    print(f"  Gaps between pairs: {[f'{g:.1f}s' for g in gaps]}\n")

    # Step 1: Stop running apps
    print("[1/5] Stopping any running app...")
    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)

    # Step 2: Start mic recording
    print(f"[2/5] Starting mic recording ({MIC_DURATION}s)...")
    mic_start = time.monotonic()
    mic_data = sd.rec(
        int(MIC_DURATION * LAPTOP_SR),
        samplerate=LAPTOP_SR, channels=1, dtype="float32",
    )

    # Step 3: Start robot
    time.sleep(0.3)
    print("[3/5] Starting robot (beeps + collisions)...")
    proc = start_robot(args.host)

    # Read stdout, capture MARK_START
    mark_start = None
    print("\n--- Robot output ---")
    for line in iter(proc.stdout.readline, ""):
        line = line.rstrip()
        if not line:
            continue
        laptop_time = time.monotonic()
        print(f"  {line}")
        if "MARK_START" in line:
            mark_start = laptop_time
    proc.wait()
    print("--- End robot output ---")

    sd.wait()
    captured = mic_data.flatten()
    print(f"\n  Mic recording done")

    if mark_start is None:
        print("\nFAILED: Never received MARK_START")
        return 1

    robot_offset = mark_start - mic_start
    print(f"  MARK_START at mic_t={robot_offset:.3f}s")

    # Save mic audio
    mic_path = Path("tests/beep_collision_mic.wav")
    sf.write(str(mic_path), captured, LAPTOP_SR)
    print(f"  Saved mic to {mic_path}")

    # Step 4: Fetch robot data
    print("\n[4/5] Fetching robot data...")
    local_results = Path("tests/beep_collision_positions.json")
    scp_from_robot(REMOTE_RESULTS, local_results, args.host)
    with open(local_results) as f:
        robot_data = json.load(f)

    # Step 5: Analyze
    print("\n[5/5] Analyzing...")

    # Detect beeps (tonal, bandpass around 2kHz)
    detected_beeps = detect_beep_onsets(
        captured, LAPTOP_SR, freq=BEEP_FREQ, bandwidth=150.0, threshold_db=-12.0,
        min_separation=1.0,  # beeps are ≥1.3s apart
    )
    print(f"  Detected {len(detected_beeps)} beeps at: "
          f"{[f'{t:.3f}' for t in detected_beeps]}")

    # Detect collisions (impulsive, highpass >2kHz)
    detected_collisions = detect_transient_onsets(
        captured, LAPTOP_SR, highpass_freq=3000.0,
    )
    print(f"  Detected {len(detected_collisions)} collisions at: "
          f"{[f'{t:.3f}' for t in detected_collisions]}")

    # Match beep-collision pairs
    # For each detected beep, find the nearest collision ~1s later
    print(f"\n{'='*60}")
    print("Beep → Collision Interval Analysis")
    print(f"  (Expected interval: {BEEP_COLLISION_OFFSET*1000:.0f}ms)")
    print(f"{'='*60}")

    pairs = []
    for i, bt in enumerate(detected_beeps):
        # Look for a collision between 0.5s and 2.0s after the beep
        candidates = [ct for ct in detected_collisions if 0.5 < (ct - bt) < 2.0]
        if not candidates:
            print(f"  Beep {i+1} at {bt:.3f}s: NO COLLISION FOUND in [+0.5, +2.0]s window")
            continue

        nearest = min(candidates, key=lambda ct: abs((ct - bt) - BEEP_COLLISION_OFFSET))
        interval_ms = (nearest - bt) * 1000
        error_ms = interval_ms - BEEP_COLLISION_OFFSET * 1000
        pairs.append({
            "beep_mic_t": bt,
            "collision_mic_t": nearest,
            "interval_ms": interval_ms,
            "error_ms": error_ms,
        })
        print(f"  Pair {len(pairs)}: beep {bt:.3f}s → collision {nearest:.3f}s = "
              f"{interval_ms:.0f}ms (error {error_ms:+.0f}ms)")

    if pairs:
        errors = [p["error_ms"] for p in pairs]
        intervals = [p["interval_ms"] for p in pairs]
        print(f"\n  Pairs matched: {len(pairs)}/{len(BEEP_TIMES)}")
        print(f"  Mean interval: {np.mean(intervals):.0f}ms (expected {BEEP_COLLISION_OFFSET*1000:.0f}ms)")
        print(f"  Mean error:    {np.mean(errors):+.0f}ms")
        print(f"  Std error:     {np.std(errors):.0f}ms")
        print(f"  Min/Max error: {min(errors):+.0f}ms / {max(errors):+.0f}ms")
    else:
        print(f"\n  No pairs matched!")

    # Generate plot
    plot_path = Path("tests/beep_collision_sync_plot.png")
    plot_combined(
        captured, LAPTOP_SR, mic_start, mark_start,
        robot_data, detected_beeps, detected_collisions, pairs, plot_path,
    )

    success = len(pairs) >= len(BEEP_TIMES) - 1
    print(f"\n{'='*60}")
    if success:
        print("RESULT: PASS — Beep-collision pairs detected and measured")
    else:
        print("RESULT: FAIL — Could not reliably detect pairs")
    print(f"{'='*60}\n")

    return 0 if success else 1


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