"""View a raiden robot_data.npz recording in rerun.

Loads the follower joint/gripper trajectories and logs them to rerun as
scalar timelines. Timestamps come from the npz's `timestamps` array
(int64 nanoseconds since epoch). The rerun web viewer is served on
port 9877 (default; override with --port).

Usage:
    /data/cameron/venvs/s3viz/bin/python /data/cameron/code/viz_recording.py \\
        /data/cameron/s3rec/flip_pink_cup_2026-06-30T16-52-57/robot_data.npz

By default only proprio (joints + grippers) is visualized. SVO2 video
streams live alongside on S3 but need the ZED SDK to decode — pass
--convert-svo-first with a directory of already-converted mp4s to add
image panels.
"""
import argparse
import time
from pathlib import Path

import numpy as np
import rerun as rr


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("npz_path", type=Path)
    ap.add_argument("--port", type=int, default=9877)
    ap.add_argument("--grpc-port", type=int, default=9878)
    args = ap.parse_args()

    print(f"Loading {args.npz_path}...")
    z = np.load(args.npz_path, allow_pickle=True)
    n = len(z["timestamps"])
    t_ns = z["timestamps"].astype(np.int64)
    t_s = (t_ns - t_ns[0]) / 1e9  # relative seconds since recording start
    print(f"  {n} frames spanning {t_s[-1]:.2f} s ({n / t_s[-1]:.1f} Hz)")

    rr.init("raiden_recording", spawn=False)
    server_uri = rr.serve_grpc(grpc_port=args.grpc_port)
    rr.serve_web_viewer(web_port=args.port, open_browser=False)
    print()
    print("=" * 60)
    print(f"  Viewer:  http://localhost:{args.port}"
          f"?url=rerun%2Bhttp%3A%2F%2F127.0.0.1%3A{args.grpc_port}%2Fproxy")
    print(f"  Tunnel:  ssh -L {args.port}:localhost:{args.port} "
          f"-L {args.grpc_port}:localhost:{args.grpc_port} <host>")
    print("=" * 60)
    print()

    # Log everything as static timelines. Each per-joint scalar lands under
    # a namespaced entity path so rerun's UI can group them into per-arm
    # + per-DOF plots automatically.
    print("Logging joints + grippers...")
    JOINT_NAMES = ["j0", "j1", "j2", "j3", "j4", "j5"]

    for arm in ("l", "r"):
        arm_label = "left" if arm == "l" else "right"
        pos = z[f"follower_{arm}_joint_pos"]      # (N, 6)
        vel = z[f"follower_{arm}_joint_vel"]      # (N, 6)
        eff = z[f"follower_{arm}_joint_eff"]      # (N, 6)
        grip_pos = z[f"follower_{arm}_gripper_pos"][:, 0]  # (N,)
        grip_vel = z[f"follower_{arm}_gripper_vel"][:, 0]
        grip_eff = z[f"follower_{arm}_gripper_eff"][:, 0]
        cmd_7d = z[f"follower_{arm}_joint_cmd"]   # (N, 7) — 6 joints + gripper cmd

        # Time-major loop; rr.set_time before each log so the timeline is
        # dense. Rerun collapses to a single scalar update per step per key.
        for i in range(n):
            rr.set_time("time", timestamp=float(t_s[i]))
            for j, name in enumerate(JOINT_NAMES):
                rr.log(f"joints/{arm_label}/pos/{name}", rr.Scalars(float(pos[i, j])))
                rr.log(f"joints/{arm_label}/vel/{name}", rr.Scalars(float(vel[i, j])))
                rr.log(f"joints/{arm_label}/eff/{name}", rr.Scalars(float(eff[i, j])))
                rr.log(f"joints/{arm_label}/cmd/{name}", rr.Scalars(float(cmd_7d[i, j])))
            rr.log(f"gripper/{arm_label}/pos", rr.Scalars(float(grip_pos[i])))
            rr.log(f"gripper/{arm_label}/vel", rr.Scalars(float(grip_vel[i])))
            rr.log(f"gripper/{arm_label}/eff", rr.Scalars(float(grip_eff[i])))
            # cmd_7d[6] is the gripper command
            rr.log(f"gripper/{arm_label}/cmd", rr.Scalars(float(cmd_7d[i, 6])))
        print(f"  ✓ {arm_label} arm ({n} frames)")

    print()
    print("Done logging. Viewer stays up — Ctrl-C to exit.")
    try:
        while True:
            time.sleep(60)
    except KeyboardInterrupt:
        print("  bye")


if __name__ == "__main__":
    main()