#!/usr/bin/env python3
"""
Go1 RobotLab ONNX Policy - MuJoCo Simulation Deployment.

Loads the RobotLab ONNX policy exported from RoboGauge and runs it in MuJoCo
with PD position control. The ONNX model uses a 10-frame history
(450-dim stacked-by-terms input, stateless).

Usage:
    conda activate free_dog_sdk
    MUJOCO_GL=glfw mjpython deploy_go1_onnx_mujoco_lab.py
    MUJOCO_GL=glfw mjpython deploy_go1_onnx_mujoco_lab.py --terrain terrains/stairs/stairs_6.xml

Controls (matching RoboGauge keyboard convention):
    ↑ / ↓         forward / backward
    ← / →         yaw left / right
    , / .         strafe left / right
    K             stop
    R             reset robot
    Esc           quit

Architecture:
    ONNX input:  obs [1, 450]  - 10 frames x 45 dims, stacked by TERMS
    ONNX output: actions [1, 12]
    Control:     target_q = default_q + 0.25 * action, PD with Kp=28, Kd=0.7
"""

import argparse
import json
import os
import re
import signal
import sys
import time
from collections import deque
from datetime import datetime
from pathlib import Path

import mujoco
import numpy as np
import onnxruntime as ort
from mujoco import viewer

# ── path setup ──
SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_ONNX = str(SCRIPT_DIR / "policy_robotlab_15000.onnx")
ROBOT_XML = str(SCRIPT_DIR / "go1.xml")
TERRAINS_DIR = SCRIPT_DIR / "terrains"

# ── policy constants (RoboGauge Go1Config) ──
NUM_OBS = 45
NUM_ACTIONS = 12
HISTORY_LEN = 10         # RobotLab ONNX history frames
ONNX_INPUT_DIM = 450     # 45 x 10, stacked by terms

ACTION_SCALE = 0.25
KP = 28.0
KD = 0.7
CLIP_ACTIONS = 100.0
CLIP_OBS = 100.0

ANG_VEL_SCALE = 0.25
CMD_SCALE = np.array([1.0, 1.0, 1.0], dtype=np.float32)

MAX_LIN_VEL_X = 1.0
MAX_LIN_VEL_Y = 0.5
MAX_ANG_VEL = 1.0

DEFAULT_DOF_POS = np.array([
    -0.1,  0.8, -1.5,   # FR_hip, FR_thigh, FR_calf
     0.1,  0.8, -1.5,   # FL_hip, FL_thigh, FL_calf
    -0.1,  1.0, -1.5,   # RR_hip, RR_thigh, RR_calf
     0.1,  1.0, -1.5,   # RL_hip, RL_thigh, RL_calf
], dtype=np.float32)

# ── exit flag ──
EXIT = False

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

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


# ═══════════════════════════════════════════════════════════════
# XML builder — merge terrain + robot into one MuJoCo scene
# ═══════════════════════════════════════════════════════════════

def _xml_inner(text, tag):
    """Extract inner content of the first <tag>...</tag> in text."""
    m = re.search(rf"<{tag}>(.*?)</{tag}>", text, re.DOTALL)
    return m.group(1).strip() if m else ""


def build_scene_xml(robot_xml_path, terrain_xml_path=None):
    """Merge robot XML with optional terrain XML into a single scene model.

    The robot XML (go1.xml) has base_link as the root body with no free joint.
    We add a free joint and optionally merge terrain worldbody/asset elements.
    """
    robot = Path(robot_xml_path).read_text()

    # 1) Add free joint to base_link
    robot = robot.replace(
        '<body name="base_link" pos="0 0 0.35">',
        '<body name="base_link" pos="0 0 0.35">\n            <joint type="free"/>'
    )

    if terrain_xml_path is None:
        # No terrain — add a simple flat floor
        floor = (
            '\n        <geom name="floor" size="0 0 0.05" type="plane" '
            'rgba="0.5 0.9 0.9 0.1"/>\n'
        )
        robot = robot.replace(
            '<body name="base_link"',
            floor + '        <body name="base_link"'
        )
        return robot

    terrain_text = Path(terrain_xml_path).read_text()

    # 2) Merge terrain <asset> into robot <asset>
    terrain_assets = _xml_inner(terrain_text, "asset")
    if terrain_assets:
        # Insert before closing </asset> of robot (or before <worldbody> if no asset)
        robot = robot.replace('</asset>', '\n' + terrain_assets + '\n    </asset>', 1)

    # 3) Merge terrain <visual> settings
    terrain_visual = _xml_inner(terrain_text, "visual")
    if terrain_visual:
        robot = robot.replace('</visual>', '\n' + terrain_visual + '\n    </visual>', 1)

    # 4) Merge terrain worldbody elements (lights, geoms, etc.) before base_link
    terrain_wb = _xml_inner(terrain_text, "worldbody")
    if terrain_wb:
        # Remove <body> elements from terrain worldbody (we only want geoms/lights/cameras)
        terrain_wb_no_bodies = re.sub(r'<body\b.*?</body>', '', terrain_wb, flags=re.DOTALL)
        robot = robot.replace(
            '<body name="base_link"',
            terrain_wb_no_bodies.strip() + '\n        <body name="base_link"'
        )

    return robot


# ═══════════════════════════════════════════════════════════════
# Observation building
# ═══════════════════════════════════════════════════════════════

def get_projected_gravity(quat_wxyz):
    qw, qx, qy, qz = quat_wxyz
    g = np.zeros(3, dtype=np.float32)
    g[0] = 2.0 * (-qz * qx + qw * qy)
    g[1] = -2.0 * (qz * qy + qw * qx)
    g[2] = 1.0 - 2.0 * (qw * qw + qz * qz)
    return g


def quat_rotate_inverse(q_wxyz, v):
    """Rotate vector v from world frame into body frame using quaternion q."""
    q_w = q_wxyz[0]
    q_vec = np.array(q_wxyz[1:], dtype=np.float32)
    v = np.array(v, dtype=np.float32)
    a = v * (2.0 * q_w * q_w - 1.0)
    b = np.cross(q_vec, v) * q_w * 2.0
    c = q_vec * np.dot(q_vec, v) * 2.0
    return a - b + c


def quat_to_rpy_deg(q_wxyz):
    """Convert MuJoCo [w, x, y, z] quaternion to roll/pitch/yaw in degrees."""
    qw, qx, qy, qz = [float(x) for x in q_wxyz]
    sinr_cosp = 2.0 * (qw * qx + qy * qz)
    cosr_cosp = 1.0 - 2.0 * (qx * qx + qy * qy)
    roll = np.arctan2(sinr_cosp, cosr_cosp)

    sinp = 2.0 * (qw * qy - qz * qx)
    pitch = np.arcsin(np.clip(sinp, -1.0, 1.0))

    siny_cosp = 2.0 * (qw * qz + qx * qy)
    cosy_cosp = 1.0 - 2.0 * (qy * qy + qz * qz)
    yaw = np.arctan2(siny_cosp, cosy_cosp)
    return np.degrees(np.array([roll, pitch, yaw], dtype=np.float32))


def read_sensor(model, data, name, expected_dim):
    """Read a MuJoCo sensor by name, returning None if it is not present."""
    sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name)
    if sid < 0:
        return None
    adr = int(model.sensor_adr[sid])
    dim = int(model.sensor_dim[sid])
    if dim != expected_dim:
        raise ValueError(f"sensor {name!r} has dim {dim}, expected {expected_dim}")
    return np.asarray(data.sensordata[adr:adr + dim], dtype=np.float32).copy()


class ObsBuilder:
    """Build 45-dim single-frame obs and stack 10-frame history into 450-dim ONNX input.

    ONNX expects observation terms stacked by groups (not by frames):
        [ang_vel(t-9..t), gravity(t-9..t), cmd(t-9..t),
         dof_pos(t-9..t), dof_vel(t-9..t), last_action(t-9..t)]
    Within each group, frames go from oldest (t-9) to newest (t).
    """

    def __init__(self):
        self.history: deque = deque(maxlen=HISTORY_LEN)

    def reset(self):
        self.history.clear()

    def build_single_obs(self, base_ang_vel_body, base_quat_wxyz, cmd, q, dq, last_action):
        """Build 45-dim single-frame observation (RoboGauge layout)."""
        obs = np.zeros(NUM_OBS, dtype=np.float32)
        obs[0:3] = np.asarray(base_ang_vel_body, dtype=np.float32) * ANG_VEL_SCALE
        obs[3:6] = get_projected_gravity(np.asarray(base_quat_wxyz, dtype=np.float32))
        obs[6:9] = np.asarray(cmd, dtype=np.float32) * CMD_SCALE
        obs[9:21] = np.asarray(q, dtype=np.float32) - DEFAULT_DOF_POS
        obs[21:33] = np.asarray(dq, dtype=np.float32) * 0.05
        obs[33:45] = np.asarray(last_action, dtype=np.float32)
        return np.clip(obs, -CLIP_OBS, CLIP_OBS)

    def build_onnx_input(self, obs_single):
        """Stack 10-frame history into 450-dim ONNX input (by-terms format)."""
        self.history.append(obs_single.copy())

        frames = list(self.history)
        while len(frames) < HISTORY_LEN:
            frames.insert(0, np.zeros(NUM_OBS, dtype=np.float32))

        # term_dims: [ang_vel(3), gravity(3), cmd(3), dof_pos(12), dof_vel(12), last_action(12)]
        term_dims = [3, 3, 3, 12, 12, 12]

        stacked = []
        offset = 0
        for dim in term_dims:
            for f_idx in range(HISTORY_LEN):
                stacked.append(frames[f_idx][offset:offset + dim])
            offset += dim

        return np.concatenate(stacked, dtype=np.float32).reshape(1, -1)


# ═══════════════════════════════════════════════════════════════
# Keyboard input (pynput)
# ═══════════════════════════════════════════════════════════════

class KbReader:
    """Non-blocking keyboard reader using a background pynput listener.

    Key names match RoboGauge convention:
        key.up / key.down  → forward/back
        key.left / key.right → yaw
        , / .              → strafe left/right
        k                  → stop
    """

    def __init__(self):
        import queue
        import threading
        self._q = queue.Queue()
        self.held = set()
        self._running = True
        self._listener = None
        self._thread = None

    @staticmethod
    def _name(key):
        try:
            if hasattr(key, 'char') and key.char is not None:
                return key.char.lower()
        except Exception:
            pass
        return str(key).lower()

    def _worker(self):
        while self._running:
            try:
                evt, key = self._q.get(timeout=0.05)
            except Exception:
                continue
            name = self._name(key)
            if evt == 'press':
                self.held.add(name)
            else:
                self.held.discard(name)

    def start(self):
        from pynput import keyboard
        self._listener = keyboard.Listener(
            on_press=lambda k: self._q.put(('press', k)),
            on_release=lambda k: self._q.put(('release', k)),
        )
        self._listener.start()
        import threading
        self._thread = threading.Thread(target=self._worker, daemon=True)
        self._thread.start()
        print("[INFO] keyboard started (↑↓←→ , . K R Esc)")

    def is_held(self, key):
        return self._name(key) in self.held

    def snapshot(self):
        return set(self.held)

    def stop(self):
        self._running = False
        if self._listener:
            self._listener.stop()


def get_command(held):
    """Extract velocity command from held key set.

    Returns (vx, vy, yaw_rate).
    """
    if 'k' in held:
        return 0.0, 0.0, 0.0

    def axis(pos_key, neg_key, limit):
        p = pos_key in held
        n = neg_key in held
        if p == n:
            return 0.0
        return limit if p else -limit

    vx = axis('key.up', 'key.down', MAX_LIN_VEL_X)
    vy = axis(',', '.', MAX_LIN_VEL_Y)
    yaw = axis('key.left', 'key.right', MAX_ANG_VEL)
    return vx, vy, yaw


class SampleLogger:
    def __init__(self, log_dir, terrain_path, cmd, args):
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.run_dir = Path(log_dir).expanduser().resolve() / f"mujoco_stairs_sample_{ts}"
        self.run_dir.mkdir(parents=True, exist_ok=True)
        self.fp = open(self.run_dir / "steps.jsonl", "a", encoding="utf-8", buffering=1)
        meta = {
            "created_at": ts,
            "terrain": str(terrain_path) if terrain_path else "flat",
            "sample_seconds": args.sample_seconds,
            "sample_cmd": list(cmd),
            "spawn": [args.spawn_x, args.spawn_y, args.spawn_z],
            "action_scale": ACTION_SCALE,
            "clip_actions": args.clip_actions,
            "max_target_step": args.max_target_step,
            "kp": KP,
            "kd": KD,
            "control_dt": 0.02,
            "sim_dt": 0.002,
            "num_obs": NUM_OBS,
            "history_len": HISTORY_LEN,
            "onnx_input_dim": ONNX_INPUT_DIM,
        }
        (self.run_dir / "metadata.json").write_text(
            json.dumps(meta, indent=2, ensure_ascii=False))
        print(f"[INFO] sample log dir: {self.run_dir}")

    def write(self, rec):
        self.fp.write(json.dumps(rec, ensure_ascii=False) + "\n")

    def close(self):
        self.fp.flush()
        self.fp.close()


def summarize_samples(samples):
    if not samples:
        return {}

    def max_abs_field(field):
        return float(max(max(abs(x) for x in s[field]) for s in samples))

    def p99_abs_field(field):
        vals = sorted(abs(x) for s in samples for x in s[field])
        if not vals:
            return 0.0
        idx = min(len(vals) - 1, int(0.99 * (len(vals) - 1)))
        return float(vals[idx])

    rpy = np.array([s["rpy_deg"] for s in samples], dtype=np.float32)
    xyz = np.array([s["base_pos"] for s in samples], dtype=np.float32)
    return {
        "num_samples": len(samples),
        "duration": float(samples[-1]["time_sim"] - samples[0]["time_sim"]),
        "base_x_start": float(xyz[0, 0]),
        "base_x_end": float(xyz[-1, 0]),
        "base_x_progress": float(xyz[-1, 0] - xyz[0, 0]),
        "base_z_min": float(np.min(xyz[:, 2])),
        "base_z_max": float(np.max(xyz[:, 2])),
        "rpy_min_deg": np.min(rpy, axis=0).astype(float).tolist(),
        "rpy_max_deg": np.max(rpy, axis=0).astype(float).tolist(),
        "action_raw_maxabs": max_abs_field("action_raw"),
        "action_raw_p99abs": p99_abs_field("action_raw"),
        "action_applied_maxabs": max_abs_field("action"),
        "action_applied_p99abs": p99_abs_field("action"),
        "target_offset_maxabs": max_abs_field("target_offset"),
        "target_offset_p99abs": p99_abs_field("target_offset"),
        "dof_vel_maxabs": max_abs_field("dof_vel"),
        "torque_maxabs": max_abs_field("torques"),
        "fallen": bool(any(s["fallen"] for s in samples)),
    }


# ═══════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(description="Go1 RobotLab ONNX MuJoCo Deployment")
    parser.add_argument("--onnx", type=str, default=DEFAULT_ONNX, help="ONNX model path")
    parser.add_argument("--terrain", type=str, default=None,
                        help="Terrain XML path (e.g. terrains/stairs/stairs_6.xml)")
    parser.add_argument("--sample", action="store_true",
                        help="Run a non-interactive fixed-command sampling rollout")
    parser.add_argument("--sample-stairs-level", type=int, default=None,
                        help="Shortcut for --sample --terrain terrains/stairs/stairs_LEVEL.xml")
    parser.add_argument("--sample-seconds", type=float, default=20.0)
    parser.add_argument("--sample-cmd-x", type=float, default=MAX_LIN_VEL_X)
    parser.add_argument("--sample-cmd-y", type=float, default=0.0)
    parser.add_argument("--sample-cmd-yaw", type=float, default=0.0)
    parser.add_argument("--sample-log-dir", type=str, default=str(SCRIPT_DIR / "mujoco_logs"))
    parser.add_argument("--spawn-x", type=float, default=None)
    parser.add_argument("--spawn-y", type=float, default=None)
    parser.add_argument("--spawn-z", type=float, default=None)
    parser.add_argument("--clip-actions", type=float, default=CLIP_ACTIONS,
                        help="Clip ONNX actions before applying action_scale")
    parser.add_argument("--max-target-step", type=float, default=0.0,
                        help="Optional per-control-step target q slew limit in rad")
    args = parser.parse_args()

    sample_mode = args.sample or args.sample_stairs_level is not None
    if args.sample_stairs_level is not None:
        args.sample = True
        args.terrain = f"terrains/stairs/stairs_{args.sample_stairs_level}.xml"

    if args.spawn_x is None:
        args.spawn_x = -0.6 if args.sample_stairs_level is not None else 0.0
    if args.spawn_y is None:
        args.spawn_y = 0.0
    if args.spawn_z is None:
        args.spawn_z = 0.34
    args.sample_log_dir = str(Path(args.sample_log_dir).expanduser().resolve())

    # Resolve terrain path
    terrain_path = None
    if args.terrain:
        terrain_path = Path(args.terrain)
        if not terrain_path.is_absolute():
            terrain_path = SCRIPT_DIR / terrain_path
        if not terrain_path.exists():
            print(f"[ERROR] terrain not found: {terrain_path}")
            return 1

    # ── Build scene XML ──
    print(f"[INFO] robot: {ROBOT_XML}")
    print(f"[INFO] terrain: {terrain_path or 'flat floor'}")
    scene_xml = build_scene_xml(ROBOT_XML, str(terrain_path) if terrain_path else None)

    # ── Load MuJoCo model ──
    os.chdir(str(SCRIPT_DIR))  # meshdir="assets" is relative to go1.xml
    model = mujoco.MjModel.from_xml_string(scene_xml)
    data = mujoco.MjData(model)
    model.opt.timestep = 0.002
    print(f"[INFO] model: {model.nbody} bodies, qpos={model.nq}, actuators={model.nu}")
    print(f"[INFO] timestep: {model.opt.timestep}s")
    has_imu_gyro = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "Body_Gyro") >= 0
    has_imu_quat = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "Body_Quat") >= 0
    print(f"[INFO] IMU sensors: Body_Gyro={has_imu_gyro}, Body_Quat={has_imu_quat}")

    # ── Load ONNX ──
    session = ort.InferenceSession(args.onnx, providers=['CPUExecutionProvider'])
    inp = session.get_inputs()[0]
    print(f"[INFO] ONNX  input : {inp.name} {inp.shape}")
    for o in session.get_outputs():
        print(f"[INFO] ONNX  output: {o.name} {o.shape}")

    # ── Init state ──
    data.qpos[0:3] = np.array([args.spawn_x, args.spawn_y, args.spawn_z], dtype=np.float64)
    data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
    data.qpos[7:19] = DEFAULT_DOF_POS.astype(np.float64)
    data.qvel[:] = 0.0
    data.ctrl[:] = 0.0
    mujoco.mj_forward(model, data)

    # ── Control state ──
    ctrl_dt = 0.02       # 50 Hz
    sim_dt = model.opt.timestep
    steps_per_inference = max(1, int(ctrl_dt / sim_dt))
    print(f"[INFO] control: {ctrl_dt}s ({1/ctrl_dt:.0f}Hz), "
          f"sim steps per inference: {steps_per_inference}")

    last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
    action_raw = np.zeros(NUM_ACTIONS, dtype=np.float32)
    action = np.zeros(NUM_ACTIONS, dtype=np.float32)
    target_pos = DEFAULT_DOF_POS.copy()
    prev_target_pos = DEFAULT_DOF_POS.copy()
    obs_builder = ObsBuilder()
    step_count = 0
    inference_step = 0

    if sample_mode:
        cmd = np.array([
            args.sample_cmd_x,
            args.sample_cmd_y,
            args.sample_cmd_yaw,
        ], dtype=np.float32)
        cmd = np.clip(cmd, [-MAX_LIN_VEL_X, -MAX_LIN_VEL_Y, -MAX_ANG_VEL],
                      [MAX_LIN_VEL_X, MAX_LIN_VEL_Y, MAX_ANG_VEL])
        logger = SampleLogger(args.sample_log_dir, terrain_path, cmd, args)
        samples = []
        max_sim_steps = int(args.sample_seconds / sim_dt)
        print(
            f"[INFO] sampling: seconds={args.sample_seconds:.1f} "
            f"cmd=({cmd[0]:.2f},{cmd[1]:.2f},{cmd[2]:.2f}) "
            f"spawn=({args.spawn_x:.2f},{args.spawn_y:.2f},{args.spawn_z:.2f})"
        )

        while step_count < max_sim_steps and not EXIT:
            if inference_step == 0:
                imu_quat = read_sensor(model, data, "Body_Quat", 4)
                imu_ang_vel = read_sensor(model, data, "Body_Gyro", 3)
                quat_wxyz = imu_quat if imu_quat is not None else data.qpos[3:7].copy()
                if imu_ang_vel is not None:
                    base_ang_vel_body = imu_ang_vel
                else:
                    world_ang_vel = data.qvel[3:6].copy()
                    base_ang_vel_body = quat_rotate_inverse(quat_wxyz, world_ang_vel)

                q = data.qpos[7:19].copy()
                dq = data.qvel[6:18].copy()
                obs_single = obs_builder.build_single_obs(
                    base_ang_vel_body, quat_wxyz, cmd, q, dq, last_action)
                onnx_input = obs_builder.build_onnx_input(obs_single)
                outputs = session.run(None, {'obs': onnx_input})
                action_raw = outputs[0][0].astype(np.float32)
                action = np.clip(action_raw, -args.clip_actions, args.clip_actions)
                last_action = action.copy()

                target_pos_raw = DEFAULT_DOF_POS + action * ACTION_SCALE
                if args.max_target_step > 0.0:
                    delta = np.clip(
                        target_pos_raw - prev_target_pos,
                        -args.max_target_step,
                        args.max_target_step,
                    )
                    target_pos = prev_target_pos + delta
                else:
                    target_pos = target_pos_raw
                prev_target_pos = target_pos.copy()

            current_pos = data.qpos[7:19]
            current_vel = data.qvel[6:18]
            torques = KP * (target_pos - current_pos) - KD * current_vel
            torques = np.clip(torques, -33.5, 33.5)
            data.ctrl[:] = torques.astype(np.float64)

            mujoco.mj_step(model, data)

            if inference_step == 0:
                quat_wxyz = data.qpos[3:7].copy()
                rpy_deg = quat_to_rpy_deg(quat_wxyz)
                fallen = bool(
                    data.qpos[2] < 0.16
                    or abs(rpy_deg[0]) > 60.0
                    or abs(rpy_deg[1]) > 60.0
                    or not np.all(np.isfinite(data.qpos))
                )
                rec = {
                    "step": int(step_count),
                    "control_step": int(step_count // steps_per_inference),
                    "time_sim": float(data.time),
                    "cmd": cmd.astype(float).tolist(),
                    "base_pos": data.qpos[0:3].astype(float).tolist(),
                    "base_quat": quat_wxyz.astype(float).tolist(),
                    "rpy_deg": rpy_deg.astype(float).tolist(),
                    "base_lin_vel": data.qvel[0:3].astype(float).tolist(),
                    "base_ang_vel": data.qvel[3:6].astype(float).tolist(),
                    "dof_pos": current_pos.astype(float).tolist(),
                    "dof_vel": current_vel.astype(float).tolist(),
                    "action_raw": action_raw.astype(float).tolist(),
                    "action": action.astype(float).tolist(),
                    "target_pos": target_pos.astype(float).tolist(),
                    "target_offset": (target_pos - DEFAULT_DOF_POS).astype(float).tolist(),
                    "torques": torques.astype(float).tolist(),
                    "fallen": fallen,
                }
                logger.write(rec)
                samples.append(rec)
                if rec["control_step"] % 50 == 0:
                    print(
                        f"[sample {rec['control_step']:04d}] "
                        f"t={rec['time_sim']:.2f} x={rec['base_pos'][0]:.2f} "
                        f"z={rec['base_pos'][2]:.2f} rpy={np.round(rpy_deg, 1)} "
                        f"act_max={np.max(np.abs(action)):.2f} "
                        f"target_off={np.max(np.abs(target_pos - DEFAULT_DOF_POS)):.2f}"
                    )
                if fallen:
                    print(f"[WARN] sample stopped: fallen at t={data.time:.2f}s")
                    break

            step_count += 1
            inference_step = (inference_step + 1) % steps_per_inference

        summary = summarize_samples(samples)
        (logger.run_dir / "summary.json").write_text(
            json.dumps(summary, indent=2, ensure_ascii=False))
        logger.close()
        print("[INFO] sample summary:")
        print(json.dumps(summary, indent=2, ensure_ascii=False))
        print(f"[INFO] sample saved: {logger.run_dir}")
        return 0

    # ── Keyboard ──
    kb = KbReader()
    kb.start()

    # ── Viewer ──
    view = viewer.launch_passive(model, data)
    # Track the robot body
    body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
    if body_id >= 0:
        view.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
        view.cam.trackbodyid = body_id
        view.cam.distance = 2.5
        view.cam.elevation = -20
        view.cam.azimuth = 60
    print("[INFO] viewer launched")

    loop_start = time.time()

    while view.is_running() and not EXIT:
        held = kb.snapshot()

        # ── Quit ──
        if 'key.esc' in held:
            break

        # ── Command from keyboard ──
        vx, vy, yaw = get_command(held)
        cmd = np.array([vx, vy, yaw], dtype=np.float32)

        # ── Reset ──
        if 'r' in held:
            data.qpos[0:3] = np.array([0.0, 0.0, 0.34], dtype=np.float64)
            data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
            data.qpos[7:19] = DEFAULT_DOF_POS.astype(np.float64)
            data.qvel[:] = 0.0
            data.ctrl[:] = 0.0
            last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
            action_raw = np.zeros(NUM_ACTIONS, dtype=np.float32)
            action = np.zeros(NUM_ACTIONS, dtype=np.float32)
            target_pos = DEFAULT_DOF_POS.copy()
            prev_target_pos = DEFAULT_DOF_POS.copy()
            obs_builder.reset()
            mujoco.mj_forward(model, data)
            print("[RESET]")

        # ── Inference ──
        if inference_step == 0:
            # Match RoboGauge: policy obs uses XML IMU gyro and framequat sensors.
            # Fall back to qpos/qvel only for XMLs without those sensors.
            imu_quat = read_sensor(model, data, "Body_Quat", 4)
            imu_ang_vel = read_sensor(model, data, "Body_Gyro", 3)
            quat_wxyz = imu_quat if imu_quat is not None else data.qpos[3:7].copy()
            if imu_ang_vel is not None:
                base_ang_vel_body = imu_ang_vel
            else:
                world_ang_vel = data.qvel[3:6].copy()
                base_ang_vel_body = quat_rotate_inverse(quat_wxyz, world_ang_vel)

            # Joint state (qpos[7:19] is FR,FL,RR,RL — matches policy order)
            q = data.qpos[7:19].copy()
            dq = data.qvel[6:18].copy()

            # Build obs
            obs_single = obs_builder.build_single_obs(
                base_ang_vel_body, quat_wxyz, cmd, q, dq, last_action)
            onnx_input = obs_builder.build_onnx_input(obs_single)

            # Run ONNX
            outputs = session.run(None, {'obs': onnx_input})
            action_raw = outputs[0][0].astype(np.float32)  # [12]
            action = np.clip(action_raw, -args.clip_actions, args.clip_actions)
            last_action = action.copy()

            target_pos_raw = DEFAULT_DOF_POS + action * ACTION_SCALE
            if args.max_target_step > 0.0:
                delta = np.clip(
                    target_pos_raw - prev_target_pos,
                    -args.max_target_step,
                    args.max_target_step,
                )
                target_pos = prev_target_pos + delta
            else:
                target_pos = target_pos_raw
            prev_target_pos = target_pos.copy()

        # ── PD control ──
        current_pos = data.qpos[7:19]
        current_vel = data.qvel[6:18]
        torques = KP * (target_pos - current_pos) - KD * current_vel
        torques = np.clip(torques, -33.5, 33.5)
        data.ctrl[:] = torques.astype(np.float64)

        mujoco.mj_step(model, data)
        view.sync()

        # Real-time sync — use sim_dt because step_count increments every sim step
        expected_time = step_count * sim_dt
        elapsed = time.time() - loop_start
        if 0 < expected_time - elapsed < ctrl_dt:
            time.sleep(expected_time - elapsed)

        step_count += 1
        inference_step = (inference_step + 1) % steps_per_inference

        # Periodic status
        if step_count % 200 == 0:
            z = data.qpos[2]
            lin_vel_abs = np.linalg.norm(data.qvel[0:3])
            print(f"[{step_count}] cmd=({vx:.1f},{vy:.1f},{yaw:.1f}) "
                  f"z={z:.3f} |v|={lin_vel_abs:.2f} "
                  f"act[0:4]={np.round(action[:4], 3)}")

    kb.stop()
    view.close()
    print("[INFO] done.")
    return 0


if __name__ == "__main__":
    sys.exit(main())
