#!/usr/bin/env python3
"""
deploy_onnx_pro_sdk.py

Deploy ONNX policy on Unitree Go1 PRO using go1_pro_sdk (no official SDK needed).
Works on non-EDU Go1 PRO.

Uses direct MCU UDP communication with Blowfish encryption.
Joint order matches ONNX directly: FR→FL→RR→RL (hip/thigh/calf per leg).
Observation: 45-dim, matches go1_sim2sim.py.

Setup:
  cd /path/to/go1_pro_sdk && pip install -e .
  pip install onnxruntime numpy

Before running, kill the Pi's sport processes:
  ssh pi@192.168.123.161
  sudo pkill -9 -f keep_sport_alive
  sudo pkill -9 -f Legged_sport
  sudo pkill -9 -f appTransit

ALWAYS suspend the robot for initial testing.

Usage:
  # Step 0: check 45-dim observation vector (no motors, no ONNX)
  python deploy_onnx_pro_sdk.py --obs-check --kill-sport

  # Step 1: test remote controller data (no motors, no ONNX)
  python deploy_onnx_pro_sdk.py --monitor --kill-sport

  # Step 2-5: state machine with RC (R2=go/stop, L2=emergency stop)
  python deploy_onnx_pro_sdk.py --onnx policy.onnx --kill-sport

  # Fixed commands, no state machine
  python deploy_onnx_pro_sdk.py --onnx policy.onnx --no-rc --no-sm --kill-sport
"""

import argparse
import json
import signal
import subprocess
import sys
import time
from datetime import datetime
from enum import Enum
from pathlib import Path

import numpy as np
import onnxruntime as ort

from go1_pro_sdk import (
    MCUClient, LowCmd, MotorCmd, MotorMode,
    apply_safety, JOINT_NAMES,
)

# ─── Policy constants (matches go1_sim2sim.py) ───
NUM_OBS = 57  # no-linvel: gyro(3)+gravity(3)+dof_pos(12)+dof_vel(12)+last_actions(12)+commands(3)+contacts(12)
NUM_ACTIONS = 12
ACTION_SCALE = 0.05
CLIP_ACTIONS = 23.7
CLIP_OBS = 100.0

DEFAULT_ANGLES = np.array([
    -0.0,  0.9, -1.8,   # FR
     0.0,  0.9, -1.8,   # FL
    -0.0,  0.9, -1.8,   # RR
     0.0,  0.9, -1.8,   # RL
], dtype=np.float32)

EXIT = False


def _sig_handler(signum, frame):
    global EXIT
    EXIT = True


signal.signal(signal.SIGINT, _sig_handler)
signal.signal(signal.SIGTERM, _sig_handler)


# ─── State machine states ───
class State(Enum):
    IDLE = "IDLE"              # full damping, wait for R2
    CALIBRATE = "CALIBRATE"    # ramping to default pose
    HOLD = "HOLD"              # holding default pose, wait for R2 to start RL
    RL = "RL"                  # running ONNX policy


# ─── Quaternion math ───
def quat_to_rot_matrix(q):
    w, x, y, z = q
    return np.array([
        [1 - 2*y*y - 2*z*z,     2*x*y - 2*w*z,     2*x*z + 2*w*y],
        [    2*x*y + 2*w*z, 1 - 2*x*x - 2*z*z,     2*y*z - 2*w*x],
        [    2*x*z - 2*w*y,     2*y*z + 2*w*x, 1 - 2*x*x - 2*y*y],
    ], dtype=np.float32)


def get_projected_gravity(quaternion):
    R = quat_to_rot_matrix(quaternion)
    return (R.T @ np.array([0., 0., -1.], dtype=np.float32)).astype(np.float32)


# ─── Observation ───
def compute_obs(imu, motor_states, commands, last_actions):
    """57-dim no-linvel: gyro(3)+gravity(3)+dof_pos(12)+dof_vel(12)+last_actions(12)+commands(3)+contacts(12)."""
    obs = np.zeros(NUM_OBS, dtype=np.float32)

    gx, gy, gz = imu.gyroscope
    obs[0:3] = np.array([gx, gy, gz], dtype=np.float32) * 0.25
    obs[3:6] = get_projected_gravity(imu.quaternion)

    dof_pos = np.array([motor_states[i].q for i in range(12)], dtype=np.float32)
    dof_vel = np.array([motor_states[i].dq for i in range(12)], dtype=np.float32)
    obs[6:18] = (dof_pos - DEFAULT_ANGLES) * 1.0
    obs[18:30] = dof_vel * 0.05

    obs[30:42] = last_actions
    obs[42:45] = np.array(commands, dtype=np.float32) * np.array([2.0, 2.0, 0.25], dtype=np.float32)
    obs[45:57] = 0.0  # contact forces = 0

    obs = np.clip(obs, -CLIP_OBS, CLIP_OBS)
    obs = np.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0)
    return obs


def state_ok(state):
    for i in range(12):
        if not np.isfinite(state.motorState[i].q):
            return False, f"motorState[{i}].q non-finite"
    gn = float(np.linalg.norm(get_projected_gravity(state.imu.quaternion)))
    if gn < 0.5 or gn > 1.5:
        return False, f"gravity norm={gn:.3f}"
    return True, "ok"


# ─── ONNX policy ───
class OnnxPolicy:
    def __init__(self, onnx_path):
        self.session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
        self.input_name = self.session.get_inputs()[0].name
        print(f"[INFO] ONNX: {onnx_path}")
        print(f"[INFO]   Input : {self.input_name} {self.session.get_inputs()[0].shape}")
        print(f"[INFO]   Output: {[(o.name, o.shape) for o in self.session.get_outputs()]}")

    def __call__(self, obs):
        out = self.session.run(None, {self.input_name: obs.reshape(1, -1).astype(np.float32)})[0][0]
        return np.asarray(out, dtype=np.float32)


# ─── Remote controller ───
def get_rc_commands(state, args):
    r = state.remote
    vx = r.ly * args.rc_vx_scale
    vy = -r.lx * args.rc_vy_scale
    wz = -r.rx * args.rc_wz_scale
    return np.array([vx, vy, wz], dtype=np.float32)


class RCEdgeDetector:
    """Detect rising/falling edges on RC buttons."""
    def __init__(self):
        self._prev = set()

    def update(self, state):
        current = set(state.remote.pressed)
        rising = current - self._prev
        falling = self._prev - current
        self._prev = current
        return rising, falling

    def rose(self, state, button):
        """True on rising edge of button."""
        current = set(state.remote.pressed)
        prev = self._prev
        self._prev = current
        return button not in prev and button in current


# ─── JSONL logger ───
class JsonlLogger:
    def __init__(self, log_dir, args):
        self.enabled = bool(log_dir)
        self.fp = None
        self.run_dir = None
        self.flush_every = 50
        if not self.enabled:
            return
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.run_dir = Path(log_dir).expanduser().resolve() / f"deploy_run_{ts}"
        self.run_dir.mkdir(parents=True, exist_ok=True)
        meta = {
            "created_at": ts,
            "num_obs": NUM_OBS,
            "num_actions": NUM_ACTIONS,
            "action_scale": ACTION_SCALE,
            "clip_actions": CLIP_ACTIONS,
            "clip_obs": CLIP_OBS,
            "default_angles": DEFAULT_ANGLES.tolist(),
            "joint_names": list(JOINT_NAMES),
        }
        for k, v in vars(args).items():
            if isinstance(v, (str, int, float, bool, type(None))):
                meta[k] = v
        (self.run_dir / "metadata.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False))
        self.fp = open(self.run_dir / "steps.jsonl", "a", encoding="utf-8", buffering=1)
        print(f"[INFO] Log dir: {self.run_dir}")

    def log(self, step, **kw):
        if not self.enabled:
            return
        rec = {"step": int(step), "time_wall": time.time()}
        for k, v in kw.items():
            if isinstance(v, np.ndarray):
                rec[k] = np.asarray(v, dtype=np.float32).reshape(-1).tolist()
            elif isinstance(v, (np.float32, np.float64)):
                rec[k] = float(v)
            elif isinstance(v, (np.int32, np.int64)):
                rec[k] = int(v)
            else:
                rec[k] = v
        self.fp.write(json.dumps(rec, ensure_ascii=False) + "\n")
        if step % self.flush_every == 0:
            self.fp.flush()

    def close(self):
        if self.fp:
            self.fp.flush(); self.fp.close()
            print(f"[INFO] Log saved: {self.run_dir}")


# ─── Display helpers ───
def fmt_rc(state):
    r = state.remote
    sticks = f"lx={r.lx:+.2f} ly={r.ly:+.2f} rx={r.rx:+.2f} ry={r.ry:+.2f} L2={r.L2:.2f}"
    btns = ",".join(r.pressed) if r.pressed else "none"
    return f"sticks:[{sticks}]  btns:[{btns}]"


def fmt_joints(state):
    parts = []
    for i in range(12):
        parts.append(f"{state.motorState[i].q:+6.3f}")
    return " ".join(parts)


# ─── Observation check mode ───
def run_obs_check(args):
    """Compute 45-dim observation and print each segment. No inference, no motors."""
    print("""
╔══════════════════════════════════════════════════════════════╗
║  OBS-CHECK MODE — verify 45-dim observation vector          ║
║  No inference, no motors. Compare with sim2sim reference.   ║
║                                                              ║
║  Ensure:                                                     ║
║  1. Robot SUSPENDED (sling / foot stand)                     ║
║  2. Network connected to robot (192.168.123.x)               ║
║  3. Use --kill-sport to auto-kill Pi sport processes         ║
╚══════════════════════════════════════════════════════════════╝
""")
    input("Press Enter to start obs check...")

    if args.kill_sport:
        kill_sport_processes(args.pi_host, args.pi_user)

    print("[INFO] Connecting to MCU...")
    client = MCUClient()
    logger = JsonlLogger(args.log_dir, args)

    try:
        print("[INFO] Waking MCU...")
        client.wake_mcu(n_frames=50, dt=0.01)

        state = client.recv_state(timeout=2.0)
        if state is None:
            print("[ERROR] No state received.")
            return 1

        print(f"[INFO] Connected. Battery={state.bms.SOC}%")
        print(f"[INFO] RPY: {np.round(np.degrees(state.imu.rpy), 1)} deg")
        print("[INFO] Computing 45-dim obs each step. Compare with sim2sim reference.\n")

        # Sim2sim reference (stationary, default pose, zero cmd)
        print("─" * 70)
        print("OBSERVATION LAYOUT (45-dim):")
        print("  obs[ 0: 3]  gyro * 0.25          →  ~[0, 0, 0] when stationary")
        print("  obs[ 3: 6]  projected gravity     →  ~[0, 0, -1] when upright")
        print("  obs[ 6:18]  (dof_pos-def) * 1.0   →  ~[0]*12 when at default pose")
        print("  obs[18:30]  dof_vel * 0.05         →  ~[0]*12 when stationary")
        print("  obs[30:42]  last_actions           →  [0]*12 initially")
        print("  obs[42:45]  commands*[2,2,0.25]    →  [0,0,0] with zero cmd")
        print("─" * 70)

        last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
        step = 0
        dt = 1.0 / max(1.0, args.rate_hz)
        next_t = time.perf_counter()

        while not EXIT:
            new_state = client.recv_latest()
            if new_state is not None:
                state = new_state
            if state is None:
                time.sleep(0.001)
                continue

            # Get commands from RC or fixed
            r = state.remote
            commands = np.array([r.ly * args.rc_vx_scale,
                                 -r.lx * args.rc_vy_scale,
                                 -r.rx * args.rc_wz_scale], dtype=np.float32)

            # Compute observation
            obs = compute_obs(state.imu, state.motorState, commands, last_actions)

            dof_pos = np.array([state.motorState[i].q for i in range(12)])
            dof_vel = np.array([state.motorState[i].dq for i in range(12)])
            gx, gy, gz = state.imu.gyroscope
            grav = get_projected_gravity(state.imu.quaternion)

            # Detailed print
            print(f"\n{'='*70}")
            print(f"[STEP {step}] bat={state.bms.SOC}%  {fmt_rc(state)}")
            print(f"{'='*70}")
            print(f"  obs[ 0: 3] gyro*0.25      = {np.round(obs[0:3], 4)}")
            print(f"            raw gyro (rad/s) = [{gx:+.4f}, {gy:+.4f}, {gz:+.4f}]")
            print(f"  obs[ 3: 6] proj gravity    = {np.round(obs[3:6], 4)}")
            print(f"            |gravity|         = {np.linalg.norm(grav):.4f} (expect ~1.0)")
            print(f"  obs[ 6:18] dof_pos_rel     = {np.round(obs[6:18], 3)}")
            diff_from_default = dof_pos - DEFAULT_ANGLES
            print(f"            |diff| max        = {np.max(np.abs(diff_from_default)):.4f}")
            print(f"            actual pos        = {np.round(dof_pos, 3)}")
            print(f"            default_angles    = {np.round(DEFAULT_ANGLES, 3)}")
            print(f"  obs[18:30] dof_vel*0.05    = {np.round(obs[18:30], 4)}")
            print(f"            |dof_vel| max     = {np.max(np.abs(dof_vel)):.4f}")
            print(f"  obs[30:42] last_actions    = {np.round(obs[30:42], 4)}")
            print(f"  obs[42:45] commands*scale  = {np.round(obs[42:45], 4)}")
            print(f"            raw cmd [vx,vy,wz]= {np.round(commands, 3)}")
            print(f"  obs min={obs.min():.3f}  max={obs.max():.3f}  mean={obs.mean():.3f}")

            # Warn if values look suspicious
            if np.linalg.norm(grav) < 0.5 or np.linalg.norm(grav) > 1.5:
                print(f"  ⚠ WARN: gravity norm is off! Robot might not be upright.")
            if np.max(np.abs(dof_pos)) < 1e-6:
                print(f"  ⚠ WARN: dof_pos is all zeros — state might be stale!")

            logger.log(
                step,
                mode="obs_check",
                obs=obs,
                dof_pos=dof_pos,
                dof_vel=dof_vel,
                base_ang_vel=np.array([gx, gy, gz]),
                projected_gravity=grav,
                imu_rpy_deg=np.degrees(state.imu.rpy),
                commands=commands,
                rc_buttons=r.pressed,
                battery_soc=state.bms.SOC,
            )

            step += 1

            next_t += dt
            sleep = next_t - time.perf_counter()
            if sleep > 0:
                time.sleep(sleep)
            else:
                next_t = time.perf_counter()

    finally:
        if logger is not None:
            logger.close()
        client.close()
        print("[INFO] Done.")


# ─── Monitor mode ───
def run_monitor(args):
    """Read and display RC + IMU + joint data. No motor commands sent."""
    print("""
╔══════════════════════════════════════════════════════════════╗
║  MONITOR MODE — no motor commands, RC test only             ║
║                                                              ║
║  Ensure:                                                     ║
║  1. Robot SUSPENDED (sling / foot stand)                     ║
║  2. Network connected to robot (192.168.123.x)               ║
║  3. Use --kill-sport to auto-kill Pi sport processes         ║
║     (or manually: ssh pi@192.168.123.161, pkill -9 ...)     ║
╚══════════════════════════════════════════════════════════════╝
""")
    input("Press Enter to start monitoring...")

    if args.kill_sport:
        kill_sport_processes(args.pi_host, args.pi_user)

    print("[INFO] Connecting to MCU...")
    client = MCUClient()
    logger = JsonlLogger(args.log_dir, args)

    try:
        print("[INFO] Waking MCU (damping frames)...")
        client.wake_mcu(n_frames=50, dt=0.01)

        state = client.recv_state(timeout=2.0)
        if state is None:
            print("[ERROR] No state received. Check connection and Pi processes.")
            return 1

        print(f"[INFO] Connected. Battery={state.bms.SOC}%")
        print("[INFO] Monitoring RC + IMU + joint data. Ctrl+C to exit.\n")

        edge = RCEdgeDetector()
        step = 0
        dt = 1.0 / max(1.0, args.rate_hz)
        next_t = time.perf_counter()

        while not EXIT:
            new_state = client.recv_latest()
            if new_state is not None:
                state = new_state

            if state is None:
                time.sleep(0.001)
                continue

            rising, falling = edge.update(state)

            rpy = np.degrees(state.imu.rpy)
            r = state.remote
            dof_pos = np.array([state.motorState[i].q for i in range(12)])
            dof_vel = np.array([state.motorState[i].dq for i in range(12)])

            print(f"\n─── [step {step}] bat={state.bms.SOC}% ───")
            print(f"  IMU  roll={rpy[0]:+.1f}  pitch={rpy[1]:+.1f}  yaw={rpy[2]:+.1f}")
            print(f"  RC   {fmt_rc(state)}")
            if rising:
                print(f"  RC   ↑ RISING: {sorted(rising)}")
            if falling:
                print(f"  RC   ↓ FALLING: {sorted(falling)}")
            print(f"  JOINT {fmt_joints(state)}")
            print(f"  JOINTS: {' '.join(f'{n:6s}' for n in JOINT_NAMES)}")

            logger.log(
                step,
                mode="monitor",
                dof_pos=dof_pos,
                dof_vel=dof_vel,
                base_ang_vel=np.array(state.imu.gyroscope),
                projected_gravity=get_projected_gravity(state.imu.quaternion),
                imu_rpy_deg=np.degrees(state.imu.rpy),
                rc_lx=r.lx, rc_ly=r.ly, rc_rx=r.rx, rc_ry=r.ry,
                rc_buttons=r.pressed,
                rc_rising=list(rising), rc_falling=list(falling),
                battery_soc=state.bms.SOC,
            )

            step += 1

            next_t += dt
            sleep = next_t - time.perf_counter()
            if sleep > 0:
                time.sleep(sleep)
            else:
                next_t = time.perf_counter()

    finally:
        if logger is not None:
            logger.close()
        client.close()
        print("[INFO] Done.")


# ─── Ramp to default pose ───
def ramp_to_default(client, args, state):
    """Slowly interpolate from current pose to DEFAULT_ANGLES (~2s).

    Position protect is DISABLED during ramp because the movement is inherently
    large (e.g. knee from -2.79 → -1.80, a 1.0 rad sweep). Only joint hard
    limits and torque limiting are active.
    """
    print("[INFO] Ramping to default pose (~2s)...")
    current = np.array([state.motorState[i].q for i in range(12)], dtype=np.float32)
    error = current - DEFAULT_ANGLES

    if np.max(np.abs(error)) < 0.05:
        print("[INFO] Already near default pose.")
        return state

    print(f"[INFO] Max error: {np.max(np.abs(error)):.2f} rad")
    ramp_steps = 200
    step_err = error / ramp_steps

    for i in range(ramp_steps):
        if EXIT:
            return state
        new_state = client.recv_latest()
        if new_state is not None:
            state = new_state

        targets = DEFAULT_ANGLES + (error - step_err * min(i + 1, ramp_steps))

        cmd = LowCmd()
        for j in range(12):
            cmd.set_motor(j, MotorCmd(
                mode=MotorMode.Servo,
                q=float(targets[j]), dq=0.0, tau=0.0,
                Kp=args.kp_cal, Kd=args.kd_cal,
            ))
        # Ramp: only joint limits + torque limit, NO position_protect
        # (the movement is inherently >0.5 rad for many joints)
        apply_safety(cmd, state, power_factor=args.power_factor,
                     position_limit_on=True, position_protect_limit=None)
        client.send(cmd)
        time.sleep(0.01)

        if i % 50 == 0:
            actual = np.array([state.motorState[j].q for j in range(12)])
            print(f"  ramp {i}/{ramp_steps}  "
                  f"target_err={np.max(np.abs(targets - DEFAULT_ANGLES)):.3f}  "
                  f"actual_err={np.max(np.abs(actual - DEFAULT_ANGLES)):.3f}")

    # Hold at default briefly (position_protect disabled — ramp continuation)
    for _ in range(50):
        if EXIT:
            return state
        new_state = client.recv_latest()
        if new_state is not None:
            state = new_state

        cmd = LowCmd()
        for j in range(12):
            cmd.set_motor(j, MotorCmd(
                mode=MotorMode.Servo,
                q=float(DEFAULT_ANGLES[j]), dq=0.0, tau=0.0,
                Kp=args.kp, Kd=args.kd,
            ))
        apply_safety(cmd, state, power_factor=args.power_factor,
                     position_limit_on=True, position_protect_limit=None)
        client.send(cmd)
        time.sleep(0.01)

    print("[INFO] Default pose reached.")
    return state


def send_hold_cmd(client, state, args):
    """Send servo commands holding DEFAULT_ANGLES.
    Position protect is disabled — hold is a fixed safe pose.
    """
    cmd = LowCmd()
    for j in range(12):
        cmd.set_motor(j, MotorCmd(
            mode=MotorMode.Servo,
            q=float(DEFAULT_ANGLES[j]), dq=0.0, tau=0.0,
            Kp=args.kp, Kd=args.kd,
        ))
    apply_safety(cmd, state, power_factor=args.power_factor,
                 position_limit_on=True, position_protect_limit=None)
    client.send(cmd)


def send_rl_cmd(client, state, targets, args):
    """Send servo commands for RL policy targets."""
    cmd = LowCmd()
    for j in range(12):
        cmd.set_motor(j, MotorCmd(
            mode=MotorMode.Servo,
            q=float(targets[j]), dq=0.0, tau=0.0,
            Kp=args.kp, Kd=args.kd,
        ))
    pp_limit = args.position_protect_limit if args.position_protect_limit > 0 else None
    apply_safety(cmd, state, power_factor=args.power_factor,
                 position_limit_on=True, position_protect_limit=pp_limit)
    client.send(cmd)


# ─── Startup instructions ───
STARTUP_BANNER = """
╔══════════════════════════════════════════════════════════════╗
║  Ensure:                                                     ║
║  1. Robot SUSPENDED (sling / foot stand)                     ║
║  2. Network connected to robot (192.168.123.x)               ║
║  3. Use --kill-sport to auto-kill Pi sport processes         ║
║     (or manually: ssh pi@192.168.123.161, pkill -9 ...)     ║
║                                                              ║
║  Controls: R2=go/stop, L2=emergency-stop(→IDLE)             ║
║            Left-stick=move, Right-stick=turn                 ║
╚══════════════════════════════════════════════════════════════╝
"""


# ─── State machine deploy mode ───
def run_deploy(args):
    print(STARTUP_BANNER)
    input("Press Enter when ready...")

    if args.kill_sport:
        kill_sport_processes(args.pi_host, args.pi_user)

    print("[INFO] Connecting to MCU...")
    client = MCUClient()
    logger = None

    try:
        print("[INFO] Waking MCU...")
        client.wake_mcu(n_frames=50, dt=0.01)

        state = client.recv_state(timeout=2.0)
        if state is None:
            print("[ERROR] No state received. Check connection and Pi processes.")
            return 1

        print(f"[INFO] Connected. Battery={state.bms.SOC}%")
        print(f"[INFO] RPY: {np.round(np.degrees(state.imu.rpy), 1)} deg")
        print("[INFO] Initial joint positions (rad):")
        for i in range(12):
            print(f"  {JOINT_NAMES[i]:6s}: {state.motorState[i].q:+7.3f}")

        # Load ONNX
        policy = OnnxPolicy(args.onnx)

        # Logger
        logger = JsonlLogger(args.log_dir, args)
        logger.flush_every = max(1, int(args.log_flush_every))

        # State machine
        sm_state = State.IDLE
        last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
        ema_filter = type('EMA', (), {'state': np.zeros(NUM_ACTIONS, dtype=np.float32)})()
        step = 0
        dt = 1.0 / args.rate_hz
        next_t = time.perf_counter()
        edge = RCEdgeDetector()

        # Seed edge detector with initial state
        edge.update(state)

        print(f"\n[INFO] State machine: IDLE → (R2) → CALIBRATE → HOLD → (R2) → RL")
        print(f"[INFO] Rate: {args.rate_hz}Hz. Ctrl+C to exit.\n")

        while not EXIT:
            new_state = client.recv_latest()
            if new_state is not None:
                state = new_state

            if state is None:
                time.sleep(0.001)
                continue

            # Rising edge detection (single update per tick)
            rising, falling = edge.update(state)
            r2_rose = "R2" in rising
            l2_rose = "L2" in rising
            r = state.remote

            # ── Emergency stop: L2 → IDLE (from any state) ──
            if l2_rose and sm_state != State.IDLE:
                print(f"\n[L2 EMERGENCY] {sm_state.value} → IDLE")
                sm_state = State.IDLE
                last_actions[:] = 0.0
                # Send damping
                client.send(LowCmd())  # all damping by default

            # ── State machine ──
            if sm_state == State.IDLE:
                # Full damping, no torques. Keep stream alive with damping frames.
                if step % 10 == 0:
                    client.send(LowCmd())  # all damping

                if r2_rose:
                    print("\n[R2] IDLE → CALIBRATE")
                    sm_state = State.CALIBRATE
                    state = ramp_to_default(client, args, state)
                    if EXIT:
                        break
                    sm_state = State.HOLD
                    print(f"[STATE] → HOLD (waiting for R2 to start RL)")

            elif sm_state == State.HOLD:
                send_hold_cmd(client, state, args)

                if r2_rose:
                    print("\n[R2] HOLD → RL")
                    sm_state = State.RL
                    last_actions[:] = 0.0
                    ema_filter.state[:] = 0.0

            elif sm_state == State.RL:
                if r2_rose:
                    print("\n[R2] RL → HOLD")
                    sm_state = State.HOLD
                    last_actions[:] = 0.0
                    # Ramp back to default
                    state = ramp_to_default(client, args, state)
                    if EXIT:
                        break
                    continue

                # Velocity commands from RC
                commands = get_rc_commands(state, args)

                # Observation + inference
                obs = compute_obs(state.imu, state.motorState, commands, last_actions)

                ok, reason = (True, "ok")
                if not args.no_state_check:
                    ok, reason = state_ok(state)

                if ok:
                    action_raw = policy(obs)
                    action_clipped = np.clip(action_raw, -CLIP_ACTIONS, CLIP_ACTIONS)
                    # EMA smooth to suppress policy's inherent high-freq oscillation
                    alpha = args.action_ema_alpha
                    if 0 < alpha < 1 and hasattr(ema_filter, 'state'):
                        ema_filter.state = alpha * action_clipped + (1 - alpha) * ema_filter.state
                        action_smoothed = ema_filter.state.copy()
                    else:
                        action_smoothed = action_clipped
                    last_actions = action_smoothed.copy()
                    targets = DEFAULT_ANGLES + action_smoothed * ACTION_SCALE
                else:
                    # State check failed: hold default pose
                    targets = DEFAULT_ANGLES.copy()

                if step >= args.warmup_steps:
                    send_rl_cmd(client, state, targets, args)

                # Log RL data
                logger.log(
                    step,
                    mode=sm_state.value,
                    commands=commands,
                    obs_isaac=obs,
                    action_raw=action_raw if ok else np.zeros(NUM_ACTIONS),
                    action_safe=action_clipped if ok else np.zeros(NUM_ACTIONS),
                    joint_targets=targets,
                    dof_pos=np.array([state.motorState[i].q for i in range(12)]),
                    dof_vel=np.array([state.motorState[i].dq for i in range(12)]),
                    base_ang_vel=np.array(state.imu.gyroscope),
                    projected_gravity=get_projected_gravity(state.imu.quaternion),
                    imu_rpy_deg=np.degrees(state.imu.rpy),
                    rc_lx=state.remote.lx, rc_ly=state.remote.ly,
                    rc_rx=state.remote.rx, rc_ry=state.remote.ry,
                    rc_buttons=state.remote.pressed,
                    state_ok=ok, state_reason=reason,
                )
            else:
                # Log non-RL states (minimal)
                logger.log(
                    step,
                    mode=sm_state.value,
                    dof_pos=np.array([state.motorState[i].q for i in range(12)]),
                    imu_rpy_deg=np.degrees(state.imu.rpy),
                    rc_buttons=state.remote.pressed,
                )

            # ── Status print ──
            if step % args.print_every == 0:
                rpy = np.degrees(state.imu.rpy)
                dof_pos = np.array([state.motorState[i].q for i in range(12)])
                print(f"\n[STEP {step}] state={sm_state.value} bat={state.bms.SOC}% "
                      f"rpy=[{rpy[0]:.0f},{rpy[1]:.0f},{rpy[2]:.0f}]")
                print(f"  RC: {fmt_rc(state)}")
                print(f"  joint: {np.round(dof_pos, 2)}")
                if sm_state == State.RL and 'targets' in dir():
                    print(f"  target:{np.round(targets, 2)}")

            step += 1
            if args.max_steps > 0 and step >= args.max_steps:
                print("[INFO] max_steps reached.")
                break

            # Rate limit
            next_t += dt
            sleep = next_t - time.perf_counter()
            if sleep > 0:
                time.sleep(sleep)
            else:
                next_t = time.perf_counter()

    finally:
        if logger is not None:
            logger.close()
        print("[INFO] Safe stopping...")
        client.safe_stop(n_frames=50, dt=0.002)
        client.close()
        print("[INFO] Done.")


# ─── Auto-kill Pi sport processes ───
def kill_sport_processes(host, user):
    """SSH into the Pi and kill sport-mode processes."""
    cmds = [
        "sudo pkill -9 -f keep_sport_alive",
        "sudo pkill -9 -f Legged_sport",
        "sudo pkill -9 -f appTransit",
    ]
    ssh_target = f"{user}@{host}" if user else host
    full_cmd = " && ".join(cmds)

    print(f"[INFO] Killing sport processes on {ssh_target}...")
    try:
        result = subprocess.run(
            ["ssh", ssh_target, full_cmd],
            capture_output=True, text=True, timeout=15
        )
        if result.returncode == 0:
            print("[INFO] Sport processes killed successfully.")
            return True
        else:
            stderr = result.stderr.strip()
            # pkill returns non-zero if no matching process — that's also OK
            if "no process" in stderr.lower() or not stderr:
                print("[INFO] No sport processes found (already killed).")
                return True
            print(f"[WARN] SSH returned {result.returncode}: {stderr}")
            return False
    except FileNotFoundError:
        print("[WARN] 'ssh' command not found. Please kill processes manually.")
        return False
    except subprocess.TimeoutExpired:
        print("[WARN] SSH timed out. Check network connection to Pi.")
        return False
    except Exception as e:
        print(f"[WARN] Failed to kill sport processes: {e}")
        return False


# ─── CLI ───
def main():
    parser = argparse.ArgumentParser(description="Deploy ONNX policy on Go1 PRO (no official SDK)")
    parser.add_argument("--onnx", default="better3.onnx", help="Path to ONNX model (57-dim no-linvel)")

    # Pi control
    parser.add_argument("--kill-sport", action="store_true",
                        help="Auto-kill sport processes on Pi via SSH before starting")
    parser.add_argument("--pi-host", default="192.168.123.161", help="Pi IP/hostname")
    parser.add_argument("--pi-user", default="pi", help="Pi SSH user")

    # Modes
    parser.add_argument("--obs-check", action="store_true",
                        help="Observation check mode: compute 45-dim obs, no inference, no motors")
    parser.add_argument("--monitor", action="store_true",
                        help="Monitor mode: read RC/IMU/joint data only, no motors")
    parser.add_argument("--no-rc", action="store_true",
                        help="Disable RC — use fixed --cmd-* velocities instead")
    parser.add_argument("--no-sm", action="store_true",
                        help="Disable state machine — run RL continuously")
    parser.add_argument("--dry-run", action="store_true",
                        help="Inference only, no motor commands")

    # Gains
    parser.add_argument("--kp", type=float, default=30.0)
    parser.add_argument("--kd", type=float, default=0.5)
    parser.add_argument("--kp-cal", type=float, default=20.0)
    parser.add_argument("--kd-cal", type=float, default=1.0)
    parser.add_argument("--power-factor", type=int, default=7,
                        help="Torque limit factor 1-10. tau_lim = TAU_MAX * factor/10. hip: %.1f, thigh: %.1f, knee: %.1f Nm at factor=7" % (23.7*7/10, 23.7*7/10, 35.55*7/10))

    # RC scale
    parser.add_argument("--rc-vx-scale", type=float, default=1.0)
    parser.add_argument("--rc-vy-scale", type=float, default=1.0)
    parser.add_argument("--rc-wz-scale", type=float, default=1.0)

    # Fixed commands (when --no-rc)
    parser.add_argument("--cmd-x", type=float, default=0.0)
    parser.add_argument("--cmd-y", type=float, default=0.0)
    parser.add_argument("--cmd-yaw", type=float, default=0.0)

    # Control
    parser.add_argument("--rate-hz", type=float, default=100.0)
    parser.add_argument("--warmup-steps", type=int, default=50)
    parser.add_argument("--max-steps", type=int, default=0)
    parser.add_argument("--print-every", type=int, default=50)

    # Safety
    parser.add_argument("--no-state-check", action="store_true")
    parser.add_argument("--position-protect-limit", type=float, default=1.0,
                        help="Max |target-actual| before zeroing Kp/Kd (negative=disable)")
    parser.add_argument("--action-ema-alpha", type=float, default=0.0,
                        help="EMA smoothing factor for actions (0=no filter, 1=bypass, default 0)")

    # Logging
    parser.add_argument("--log-dir", default="", help="Enable JSONL logging to this directory")
    parser.add_argument("--log-flush-every", type=int, default=50)

    args = parser.parse_args()

    if args.obs_check:
        return run_obs_check(args)
    if args.monitor:
        return run_monitor(args)
    else:
        return run_deploy(args)


if __name__ == "__main__":
    main()
