"""Record YAM follower-arm joint trajectories at high rate. No cameras, no SVO2.

Uses raiden's RobotController for hardware access. The followers can be in any
mode while recording (gravity-comp / leader-follower teleop / passive). We just
sample `get_joint_pos()` over time.

Output: a single .npz with
    joints       (N, 14) float32   — left arm (7) then right arm (7), incl. gripper
    timestamps   (N,)    float64   — seconds since recording start (t[0] = 0)
    sample_hz    float32           — measured average sample rate
    metadata     bytes              — JSON blob (notes, datetime, mode, hostname)

Usage (on the YAM machine, inside the raiden fork's venv):
    source ~/cameron/raiden_fork/.venv/bin/activate
    python ~/cameron/yam_control/yam_record.py \\
        --duration 30 --hz 100 \\
        --out ~/cameron/yam_control/recordings/$(date +%Y%m%d_%H%M%S).npz \\
        --notes "leader-teleop, reach-grasp-lift"

Optionally enable gravity-comp on the followers so you can hand-guide them:
    --gravity_comp_followers
"""
import argparse
import json
import os
import socket
import sys
import time
from datetime import datetime
from pathlib import Path

import numpy as np

# raiden lives as an editable install in the venv — just import it.
from raiden.robot.controller import RobotController


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--out", required=True, type=str,
                   help="Path to output .npz (parent dir is created).")
    p.add_argument("--duration", type=float, default=30.0,
                   help="Recording duration in seconds.")
    p.add_argument("--hz", type=float, default=100.0,
                   help="Sample rate for proprio (Hz).")
    p.add_argument("--notes", type=str, default="",
                   help="Free-text annotation saved into metadata.")
    p.add_argument("--gravity_comp_followers", action="store_true",
                   help="Put followers into gravity-comp mode (zero-stiffness) so you can "
                        "hand-guide them during recording. WARNING: hold the arms — they will "
                        "be free-floating.")
    p.add_argument("--use_right_leader", type=int, default=0,
                   help="Initialize the right leader arm (default off — pure recording).")
    p.add_argument("--use_left_leader", type=int, default=0)
    p.add_argument("--use_right_follower", type=int, default=1)
    p.add_argument("--use_left_follower", type=int, default=1)
    args = p.parse_args()

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    print("Initialising robot controller (followers only by default)...")
    rc = RobotController(
        use_right_leader=bool(args.use_right_leader),
        use_left_leader=bool(args.use_left_leader),
        use_right_follower=bool(args.use_right_follower),
        use_left_follower=bool(args.use_left_follower),
    )
    rc.check_can_interfaces()
    rc.initialize_robots(gravity_comp_mode=args.gravity_comp_followers)

    dof_per_arm = 7  # 6 arm + 1 gripper

    def read_joints14():
        ql = rc.follower_l.get_joint_pos() if rc.follower_l is not None else np.zeros(dof_per_arm)
        qr = rc.follower_r.get_joint_pos() if rc.follower_r is not None else np.zeros(dof_per_arm)
        return np.concatenate([ql, qr]).astype(np.float32)

    samples = []
    timestamps = []
    dt = 1.0 / args.hz
    print(f"\nRecording for {args.duration:.1f}s at {args.hz:.0f} Hz ...")
    print("(Move the arms now — Ctrl-C to stop early.)\n")
    t_start = time.monotonic()
    next_sample = t_start
    try:
        while True:
            now = time.monotonic()
            if now - t_start > args.duration:
                break
            if now >= next_sample:
                samples.append(read_joints14())
                timestamps.append(now - t_start)
                next_sample += dt
            else:
                time.sleep(max(0.0, next_sample - now))
    except KeyboardInterrupt:
        print("\n[Ctrl-C] stopping early")
    finally:
        # Always return to home before closing (raiden's pattern)
        try:
            rc.return_to_home()
        except Exception as e:
            print(f"return_to_home failed: {e}")

    joints = np.stack(samples, axis=0) if samples else np.zeros((0, 14), dtype=np.float32)
    ts = np.asarray(timestamps, dtype=np.float64)
    sample_hz = float((len(ts) - 1) / (ts[-1] - ts[0])) if len(ts) >= 2 else 0.0
    meta = {
        "notes": args.notes,
        "datetime": datetime.now().isoformat(timespec="seconds"),
        "hostname": socket.gethostname(),
        "duration_s": float(ts[-1] - ts[0]) if len(ts) else 0.0,
        "requested_hz": args.hz,
        "measured_hz": sample_hz,
        "n_samples": int(len(ts)),
        "use_right_follower": bool(args.use_right_follower),
        "use_left_follower": bool(args.use_left_follower),
        "gravity_comp_followers": bool(args.gravity_comp_followers),
    }
    np.savez(
        out_path,
        joints=joints,
        timestamps=ts,
        sample_hz=np.float32(sample_hz),
        metadata=np.frombuffer(json.dumps(meta).encode(), dtype=np.uint8),
    )
    print(f"\nSaved {len(joints)} samples ({sample_hz:.1f} Hz) → {out_path}")
    print(f"  duration = {meta['duration_s']:.2f}s")
    print(f"  metadata = {meta}")


if __name__ == "__main__":
    main()