tests/scripts/benchmark_python.py

"""Python equivalent of the JS streaming benchmark.

Runs the exact same motion (mirror-symmetric antennas + head Y sin)
through the Python SDK while reading state at the same nominal rate.
Output CSV matches the JS schema so the two can be compared head-to-head.

Notes on parity vs JS:
- No lead-compensation here. The JS Playback shifts antenna commands
  40 ms ahead and head 150 ms ahead. This script sends the raw motion
  at wall-clock t. We're hunting for jitter/stutter events, which show
  up the same way regardless of mean phase lag.
- The Python SDK exposes synchronous `get_current_head_pose()` /
  `get_present_antenna_joint_positions()` that read a client-side
  cache of the latest daemon state. Read timing is therefore set by
  *our* loop, not by daemon-pushed events as in JS.
"""

import csv
import math
import time

import numpy as np
from reachy_mini import ReachyMini

# Match lib/benchmark.js BENCHMARK constants exactly.
CENTER_DEG = 30.0
AMP_DEG = 20.0
PHASE1_DUR = 5.0
PHASE1_HZ = 0.5
HOLD_DUR = 2.0
PHASE3_DUR = 5.0
PHASE3_HZ = 1.5
HEAD_AMP_M = 0.03
HEAD_PHASE1_HZ = 0.3
HEAD_PHASE3_HZ = 0.5
TOTAL_DUR = PHASE1_DUR + HOLD_DUR + PHASE3_DUR
LOOP_HZ = 50.0
PERIOD = 1.0 / LOOP_HZ


def right_at(t: float) -> float:
    center = math.radians(CENTER_DEG)
    amp = math.radians(AMP_DEG)
    if t < 0:
        return center
    if t <= PHASE1_DUR:
        return center + amp * math.sin(2 * math.pi * PHASE1_HZ * t)
    if t <= PHASE1_DUR + HOLD_DUR:
        return center
    if t <= TOTAL_DUR:
        return center + amp * math.sin(
            2 * math.pi * PHASE3_HZ * (t - PHASE1_DUR - HOLD_DUR)
        )
    return center


def left_at(t: float) -> float:
    return -right_at(t)


def head_y_at(t: float) -> float:
    if t < 0:
        return 0.0
    if t <= PHASE1_DUR:
        return HEAD_AMP_M * math.sin(2 * math.pi * HEAD_PHASE1_HZ * t)
    if t <= PHASE1_DUR + HOLD_DUR:
        return 0.0
    if t <= TOTAL_DUR:
        return HEAD_AMP_M * math.sin(
            2 * math.pi * HEAD_PHASE3_HZ * (t - PHASE1_DUR - HOLD_DUR)
        )
    return 0.0


def head_pose_y(y_m: float) -> np.ndarray:
    H = np.eye(4)
    H[1, 3] = y_m
    return H


def main() -> None:
    samples: list[tuple[float, float, float, float, float, float, float]] = []
    print("Connecting to Reachy Mini…")
    with ReachyMini() as mini:
        print("Connected. Going to initial pose…")
        init_head = head_pose_y(head_y_at(0.0))
        mini.goto_target(
            init_head,
            antennas=[left_at(0.0), right_at(0.0)],
            duration=1.5,
        )
        time.sleep(0.2)

        print(f"Running benchmark: {TOTAL_DUR:.1f} s at {LOOP_HZ:.0f} Hz…")
        t0 = time.perf_counter()
        tick = 0
        while True:
            now = time.perf_counter()
            t = now - t0
            if t > TOTAL_DUR:
                break
            cmd_l = left_at(t)
            cmd_r = right_at(t)
            cmd_hy = head_y_at(t)
            mini.set_target(
                head=head_pose_y(cmd_hy),
                antennas=[cmd_l, cmd_r],
            )
            # Read the latest cached state right after sending. This is
            # the closest analogue to what the JS path measures (where
            # state events are pushed asynchronously). Read time is our
            # wall-clock t.
            try:
                actual_head = mini.get_current_head_pose()
                actual_ant = mini.get_present_antenna_joint_positions()
                samples.append((
                    t,
                    math.degrees(cmd_l),
                    math.degrees(actual_ant[0]),
                    math.degrees(cmd_r),
                    math.degrees(actual_ant[1]),
                    cmd_hy * 100.0,
                    float(actual_head[1, 3]) * 100.0,
                ))
            except Exception as e:
                print(f"  read error at t={t:.3f}: {e}")

            tick += 1
            next_tick = t0 + tick * PERIOD
            sleep_for = next_tick - time.perf_counter()
            if sleep_for > 0:
                time.sleep(sleep_for)

        print("Done. Returning to base…")
        mini.goto_target(
            np.eye(4),
            antennas=[-0.1745, 0.1745],
            duration=1.0,
        )
        time.sleep(1.1)

    stamp = time.strftime("%Y-%m-%dT%H-%M-%S")
    out = f"/Users/remi/Downloads/marionette-benchmark-python-{stamp}.csv"
    with open(out, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow([
            "t_s",
            "left_commanded_deg", "left_actual_deg",
            "right_commanded_deg", "right_actual_deg",
            "head_y_commanded_cm", "head_y_actual_cm",
        ])
        for row in samples:
            w.writerow([
                f"{row[0]:.4f}",
                f"{row[1]:.3f}", f"{row[2]:.3f}",
                f"{row[3]:.3f}", f"{row[4]:.3f}",
                f"{row[5]:.3f}", f"{row[6]:.3f}",
            ])
    print(f"Wrote {out}")
    print(f"Samples: {len(samples)}")
    if samples:
        dts = [(samples[i][0] - samples[i - 1][0]) * 1000 for i in range(1, len(samples))]
        print(
            f"State interval: mean {sum(dts)/len(dts):.1f} ms, "
            f"max {max(dts):.0f} ms, "
            f"p99 {sorted(dts)[int(len(dts)*0.99)]:.0f} ms"
        )


if __name__ == "__main__":
    main()