#!/usr/bin/env python3
"""
Arm tracking measurement — PASSIVE, commands nothing, only listens.

Run this, then start a replay in another terminal. It captures both sides of the loop:

  rt/arm_sdk    the recorder's COMMANDED q / dq / tau / kp / kd
  rt/lowstate   the robot's ACTUAL q / dq / tau_est

and separates the causes that a single "mean error" number hides:

  CLAMP-STARVED  tau feedforward pinned at TAU_CLAMP -> integral cannot supply the torque
                 the joint needs, error explodes. Fix: raise that joint's clamp.
  LAG            actual is a delayed copy of commanded -> time-shifting collapses the error.
                 Fix: feedforward / higher send rate. NOT more kp.
  DROOP          actual sits below commanded even at rest. Fix: more gravity feedforward.
  dynamic        error grows with speed but is none of the above.

  conda activate g1_env
  cd ~/Robotics_workspace/yslootahtech/G1/G1_Lootah/Manual_Recorder
  python3 arm_track_measure.py enp3s0 90          # terminal 1, then replay in terminal 2

Writes arm_track_raw.npz (re-analysable offline, no robot needed) and arm_track_report.txt
next to this script. Do NOT put these under /tmp — systemd-tmpfiles-clean wipes it.
"""
import sys, os, time, math
import numpy as np
from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_, LowState_

HERE = os.path.dirname(os.path.abspath(__file__))
RAW = os.path.join(HERE, "arm_track_raw.npz")
REPORT = os.path.join(HERE, "arm_track_report.txt")

IFACE = sys.argv[1] if len(sys.argv) > 1 else "enp3s0"
SECS = float(sys.argv[2]) if len(sys.argv) > 2 else 90.0

JN = {15: "L_sh_pitch", 16: "L_sh_roll", 17: "L_sh_yaw", 18: "L_elbow",
      19: "L_wr_roll", 20: "L_wr_pitch", 21: "L_wr_yaw",
      22: "R_sh_pitch", 23: "R_sh_roll", 24: "R_sh_yaw", 25: "R_elbow",
      26: "R_wr_roll", 27: "R_wr_pitch", 28: "R_wr_yaw"}
ARM = list(JN)

# Mirror TAU_CLAMP from g1_record_replay.py so saturation can be detected.
CLAMP = {**{j: 20.0 for j in (15, 16, 17, 18, 22, 23, 24, 25)},
         **{j: 12.0 for j in (19, 26)},
         **{j: 4.0 for j in (20, 21, 27, 28)}}

_out = open(REPORT, "w")
def P(*a):
    s = " ".join(str(x) for x in a)
    print(s); _out.write(s + "\n"); _out.flush()

ChannelFactoryInitialize(0, IFACE)
cmd, act = [], []
# Each callback appends exactly ONE tuple, so any snapshot of the list is self-consistent.
# Building separate per-field arrays while callbacks are still firing gives arrays of
# different lengths — that raced and crashed an earlier version of this script.
ChannelSubscriber("rt/arm_sdk", LowCmd_).Init(lambda m: cmd.append((
    time.time(),
    [float(m.motor_cmd[i].q) for i in ARM],
    [float(m.motor_cmd[i].dq) for i in ARM],
    [float(m.motor_cmd[i].tau) for i in ARM],
    float(m.motor_cmd[15].kp), float(m.motor_cmd[15].kd))), 50)
ChannelSubscriber("rt/lowstate", LowState_).Init(lambda m: act.append((
    time.time(),
    [float(m.motor_state[i].q) for i in ARM],
    [float(m.motor_state[i].dq) for i in ARM],
    [float(m.motor_state[i].tau_est) for i in ARM])), 50)

P(f"listening on {IFACE} for up to {SECS:.0f}s  ---  START THE REPLAY NOW")
t0 = time.time()
while not cmd and time.time() - t0 < SECS:
    time.sleep(0.1)
if not cmd:
    sys.exit("no rt/arm_sdk traffic seen - was a replay actually running?")
P("commands detected, recording until they stop\n")

last, quiet, t1 = len(cmd), 0, time.time()
while time.time() - t1 < SECS:
    time.sleep(0.5)
    quiet = quiet + 1 if len(cmd) == last else 0
    last = len(cmd)
    if quiet >= 6:                       # 3s with no new commands => replay finished
        break

C, S = list(cmd), list(act)              # freeze one snapshot each before building arrays
P(f"captured {len(C)} commands, {len(S)} states")

ct = np.array([c[0] for c in C]); cq = np.array([c[1] for c in C])
ctau = np.array([c[3] for c in C])
at = np.array([s[0] for s in S]); aq = np.array([s[1] for s in S])
adq = np.array([s[2] for s in S]); atau = np.array([s[3] for s in S])
np.savez_compressed(RAW, ct=ct, cq=cq, cdq=np.array([c[2] for c in C]), ctau=ctau,
                    kp=np.array([c[4] for c in C]), kd=np.array([c[5] for c in C]),
                    at=at, aq=aq, adq=adq, atau=atau, arm=np.array(ARM))
P(f"raw saved -> {RAW}\n")

A = np.stack([np.interp(ct, at, aq[:, j]) for j in range(14)], 1)
V = np.stack([np.interp(ct, at, adq[:, j]) for j in range(14)], 1)
err = cq - A
dt = np.median(np.diff(ct))
t = ct - ct[0]
mid = (t > 2.5) & (t < t[-1] - 4)        # exclude the ease-in and the home return

P(f"kp={C[0][4]:.0f} kd={C[0][5]:.1f}   cmd {1/dt:.1f} Hz   state {len(S)/(at[-1]-at[0]):.0f} Hz")
P(f"tau ff {'ON' if np.abs(ctau).max() > 0.01 else 'OFF'}  (max seen {np.abs(ctau).max():.2f} Nm)\n")

P("=" * 116)
P(f"{'joint':12s} {'mean':>7s} {'peak':>7s} {'@rest':>7s} {'@move':>7s} {'lag ms':>7s} "
  f"{'drop':>6s} {'%clamp':>7s} {'e|clamp':>8s} {'e|free':>7s}  verdict")
P("=" * 116)
worst = []
for k, j in enumerate(ARM):
    e = err[mid, k]; v = V[mid, k]; tau = np.abs(ctau[mid, k])
    mv, rs = np.abs(v) > 0.15, np.abs(v) < 0.02
    me, pe = math.degrees(np.abs(e).mean()), math.degrees(np.abs(e).max())
    er = math.degrees(np.abs(e[rs]).mean()) if rs.sum() > 5 else float("nan")
    em = math.degrees(np.abs(e[mv]).mean()) if mv.sum() > 5 else float("nan")
    bs, be = 0, np.abs(e).mean()
    for s in range(1, 25):
        sh = np.abs(cq[mid][:-s, k] - A[mid][s:, k]).mean()
        if sh < be: be, bs = sh, s
    drop = (1 - be / max(np.abs(e).mean(), 1e-12)) * 100
    sat = tau >= CLAMP[j] - 1e-6
    ec = math.degrees(np.abs(e[sat]).mean()) if sat.sum() > 3 else float("nan")
    ef = math.degrees(np.abs(e[~sat]).mean()) if (~sat).sum() > 3 else float("nan")
    vd = ("CLAMP-STARVED" if sat.mean() > 0.002 and ec == ec and ec > ef * 2 else
          "LAG" if drop > 50 else
          "DROOP" if (er == er and er > 0.5 and em == em and em < er * 1.5) else
          "dynamic" if (em == em and er == er and em > er * 2) else "small")
    worst.append((pe, JN[j], me, vd))
    P(f"{JN[j]:12s} {me:7.2f} {pe:7.2f} {er:7.2f} {em:7.2f} {bs*dt*1000:7.0f} "
      f"{drop:5.0f}% {sat.mean()*100:6.1f}% {ec:8.2f} {ef:7.2f}  {vd}")
P("=" * 116)
P("(mid-playback only; ease-in and home return excluded)\n")
allm = math.degrees(np.abs(err[mid]).mean())
P(f"OVERALL mean {allm:.2f} deg    worst peaks: " +
  ", ".join(f"{n} {p:.2f}d" for p, n, m, v in sorted(worst, reverse=True)[:4]))
starved = [n for p, n, m, v in worst if v == "CLAMP-STARVED"]
P("clamp-starved joints: " + (", ".join(starved) if starved else "NONE - no joint is hitting its limit"))
_out.close()