armcade

#!/usr/bin/env python

# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Cartesian end-effector keyboard teleop SIM for the Tom follower arm.

This is a *kinematic* simulation (no contact dynamics / payload). It renders the
real Tom arm — actual CAD meshes from the bundled URDF — in a meshcat browser
viewer, and lets you drive the gripper in task space with the keyboard:

    Arrow keys ........ move the gripper in the XY plane of the work surface
                        (fixed height)
    ' / . ............. move the gripper UP / DOWN along the world Z axis
    [ / ] ............. OPEN / CLOSE the gripper
    h ................. return to the home (ready) pose
    q / Esc ........... quit

The controlled point is the gripper TIP (end of the claws), solved for in XYZ by a prioritised
(null-space) IK: tip position is the PRIMARY task (a move always drives the tip where you command),
while keeping the gripper pointing down and the arm near its ready posture are SECONDARY tasks
solved only in the primary's null space — so they never fight the position. Consequence: wherever
the tip is reachable with the gripper down it stays down, and when you reach LONG (no null space
left) the position solve naturally ROTATES THE GRIPPER OUT to throw the tip further, up to a cap
(MAX_ORIENTATION_RESIDUAL_DEG) past which it's too sideways to grasp. The posture term keeps the
redundancy consistent (no drift); the wrist is left out of it so it's free to rotate out. Joint
ranges are the arms' calibrated limits (JOINT_LIMITS_DEG); verify/adjust them with the joint
sliders in --web. Gripper open/close is a separate channel (not in the IK chain).

The arm is clamped to the -X edge of the work surface (the green servo board / clamp
overhang that edge) and reaches out over the table toward +X. For the web viewer's
default camera, the arrows map: up = +X (away, out over the table), down = -X (back),
left = +Y, right = -Y.

Geometry, FK and IK all come from the SAME bundled URDF
(assets/tom.urdf) so the rendered mesh matches the solved pose.
Joint values are the same normalized degrees you send the real robot as
`<joint>.pos`, so a motion that works here maps onto hardware.

SHARED CONTROL CORE: this module is ALSO the Cartesian-control core for the real arm. The
companion runner ``-m tomctl.real`` imports it and reuses the SAME
``CartesianController`` (its null-space IK solver), the SAME keyboard/remote intent layer, and the
SAME ``load_calibration`` to drive the PHYSICAL Tom follower — only the forward-kinematics backend
differs (the sim's CAD URDF here vs the arm's vendored URDF there). Keep ``CartesianController``
backend-agnostic: it touches the robot only via ``backend.forward_kinematics`` + ``joint_names``.

Usage:
    # Build the mesh bundle once (needs trimesh; see build_sim_assets.py). The bundle is
    # committed, so you only need this to regenerate it:
    uv run --with trimesh --with lxml --with fast-simplification \
        python build_sim_assets.py

    # RECOMMENDED: browser control panel. Renders the arm with three.js in ONE page that
    # captures the keyboard directly (no terminal focus, no pynput, no OS permission; needs
    # only a browser with internet for the three.js CDN). This is the most reliable mode.
    uv run python -m tomctl.control --web

    # Live sim via meshcat + a terminal/pynput keyboard listener (keys are read in the
    # terminal, not the browser). Needs meshcat (in the `all` / `unitree_g1` extras).
    uv run --with meshcat python -m tomctl.control

    # Hands-off: play a scripted trajectory in the meshcat viewer (no keyboard needed):
    uv run --with meshcat python -m tomctl.control --demo

    # Headless self-check (no browser/keyboard, no meshcat needed): scripts each control
    # and reports whether the down-orientation and reachability hold. Good for CI.
    uv run python -m tomctl.control --check
"""

from __future__ import annotations

import argparse
import signal
import sys
import threading
import time
import xml.etree.ElementTree as ET
from dataclasses import dataclass, field
from pathlib import Path

# Joy-Con teleop (optional): load pygame BEFORE lerobot drags in OpenCV. OpenCV and pygame each vendor
# their own libSDL2; on macOS the first one loaded wins the Objective-C class registration, and if
# OpenCV's wins, pygame's joystick enumeration silently sees only SOME controllers. Importing pygame
# first keeps its SDL backend intact. Guarded on argv so keyboard-only runs don't pay for pygame.
if "--gamepad" in sys.argv:
    try:
        import pygame  # noqa: F401  (loaded first on purpose; used later via tomctl.gamepad)
    except ImportError:
        pass

import numpy as np
from numpy.typing import NDArray

from lerobot.robots.tom.kinematics import TomKinematics

from tomctl import ASSETS

DEFAULT_URDF = ASSETS / "tom.urdf"


def ik_calibration_path(robot_id: str) -> Path:
    """Per-arm IK-control calibration file: ``assets/ik_calibration_<robot_id>.json``.

    Each PHYSICAL arm has its OWN file — motor mounting, posing, and reach differ, so calibrations are
    keyed by id and NEVER shared between arms (a single shared file silently served one arm's
    calibration to the other, which looked like flipped joints). There is no shared/legacy default:
    every arm must be calibrated. This is the IK-CONTROL layer (home/rest/limits/gripper/motion),
    written by ``python -m tomctl.calibrate.arm`` and separate from (and on top of) lerobot's hardware
    (servo zero/range) calibration."""
    if not robot_id:
        raise ValueError("robot_id is required — IK calibration is strictly per-arm (no shared default).")
    return ASSETS / f"ik_calibration_{robot_id}.json"

# ----------------------------------------------------------------------------------------
# Control tuning. Step sizes are PER KEY PRESS. Holding a key gives smooth motion via OS
# key-repeat only on the pynput backend (X11 / macOS / Windows); on the terminal backend
# (Wayland / headless-SSH / untrusted macOS) keys arrive as discrete presses, so "hold to
# move" there is choppy and repeat-rate dependent — tap repeatedly instead.
# ----------------------------------------------------------------------------------------
STEP_XY_M = 0.005  # metres the COMMANDED tip target advances in the surface plane per arrow press
STEP_Z_M = 0.005  # metres the commanded target advances along world Z per up/down press
STEP_GRIPPER = 5.0  # gripper command units (0..100) per o/c press
LOOP_HZ = 50.0  # control-loop rate (Hz): slews the working target + solves IK each tick

# Motion limits for the slew from the COMMANDED (intent) target to the WORKING target the IK tracks.
# These shape every move into a smooth accel/cruise/decel profile regardless of how (or how jerkily,
# or how delayed) commands arrive — which is what keeps the real robot's motion smooth. Bounded
# Cartesian speed/accel keeps joint speeds bounded; JOINT_MAX_SPEED_DPS is a hard per-joint backstop.
CART_MAX_SPEED_MPS = 0.18  # max tip speed
CART_MAX_ACCEL_MPS2 = 1.2  # max tip acceleration (ramps speed up/down -> no jerk at start/stop)
GRIPPER_SPEED_PPS = 250.0  # gripper command units per second (0..100 travel)
# Backstop for the release-before-move gate: start a queued home/rest move after at most this long even
# if the gripper never reads fully open, so a stuck/miscalibrated gripper can't wedge the move (and, via
# is_settling(), the one-arm-at-a-time cross-arm gate) forever. Normal release is <0.5s at 250 pps.
PENDING_RELEASE_TIMEOUT_S = 2.0
JOINT_MAX_SPEED_DPS = 150.0  # per-joint angular-speed backstop (deg/s)
JOINT_MAX_ACCEL_DPS2 = 600.0  # per-joint angular-ACCEL backstop for manual jog (deg/s^2): caps how fast
# a joint's speed may change per tick so the IK's redundancy/null-space corrections can't jerk it. The
# planned pose-to-pose move is smooth via JOINT_PATH_ACCEL_DPS2; the manual jog had no joint-accel limit,
# which read as shake. High enough (reaches JOINT_MAX_SPEED in ~0.25 s) to stay responsive.

# Joint-space point-to-point moves for the named-pose path planner (home/rest/custom). These are
# trapezoidal (accel -> cruise -> decel) and synchronised across joints, so transitions are smooth.
JOINT_PATH_SPEED_DPS = 90.0  # cruise speed of a planned pose-to-pose move
JOINT_PATH_ACCEL_DPS2 = 300.0  # accel/decel for the planned move (smooth start/stop at each waypoint)

# "Facing the ground" = the tool's approach axis (its local +Z) points along world -Z. We
# constrain only this AXIS, never the heading about it. Pinning the full orientation makes
# base_yaw and the two roll joints redundant about the vertical, and the IK then drifts
# through that null space (the arm slowly twists; "right" stops meaning right) and walks
# joints into their limits until it sticks. Constraining the axis alone lets the heading
# follow naturally from where you reach, which is stable.
DOWN_AXIS = np.array([0.0, 0.0, -1.0])

# Operator JOG FRAME: the arrow keys command (forward, left) in the OPERATOR's frame, which the
# controller rotates into the robot's world XY before moving the tip. "forward" (UP) is aligned with
# the horizontal direction the arm reaches at HOME (the home tip's azimuth) so UP = out over the
# table and LEFT = the operator's left, regardless of the URDF's base-frame zero or the home's
# base_yaw. Without this, UP=world+X — which only matches "forward" if the arm happens to face +X at
# home (true for the sim, but the real arm faces ~+Y, so UP would swing the base sideways). Set this
# to a number (degrees) to override the auto heading for a different operator viewpoint.
JOG_FRAME_YAW_DEG: float | None = None
# Some arms' world XY is left-handed relative to the operator (e.g. world +X = the operator's LEFT,
# +Y = forward). A plain rotation then sends LEFT -> right. JOG_FLIP_LR mirrors the left axis so LEFT
# is the operator's left. Set per arm (calibration jog_flip_lr / runner --jog-flip-lr).
JOG_FLIP_LR: bool = False

# All six arm joints are active. With the gripper TIP as the controlled point the system is
# redundant — which would cause null-space drift — so a posture regulariser (IK_NS_POSTURE_GAIN)
# gently pulls the joints back toward the ready pose in the null space, keeping the configuration
# consistent. The wrist is left out of that pull (POSTURE_MASK) so it's free to rotate out for reach.
IK_ACTIVE_JOINTS = (
    "base_yaw",
    "shoulder_pitch",
    "shoulder_roll",
    "elbow_pitch",
    "elbow_roll",
    "wrist_pitch",
)

# CONTINUOUS joints are free-spinning (URDF type="continuous"): they have no hard stop, so their
# +-180 entries in JOINT_LIMITS_DEG are just a wrap convention, not a measured limit — calibration
# skips them. base_yaw, like the two rolls, is fully rotatable, so only the three revolute PITCH
# joints (shoulder_pitch, elbow_pitch, wrist_pitch) — and the gripper, when present — carry real
# travel stops worth calibrating.
CONTINUOUS_JOINTS = ("base_yaw", "shoulder_roll", "elbow_roll")
LIMIT_CALIBRATED_JOINTS = tuple(j for j in IK_ACTIVE_JOINTS if j not in CONTINUOUS_JOINTS)

# Per-joint (lo, hi) travel limits in degrees — the real arms' calibrated ranges (measured with
# the joint sliders). Notes: base_yaw and the two rolls are CONTINUOUS (their +-180 here is just a
# wrap convention, not a stop — see CONTINUOUS_JOINTS; _bounds/_postprocess leave them unclamped);
# the elbow's "straight" (forearm collinear) is at +90, the MIDDLE of its range — it folds to -15
# and bends back the symmetric amount to +195.
JOINT_LIMITS_DEG = {
    "base_yaw": (-180.0, 180.0),
    # shoulder_pitch widened toward the real-arm calibration (~[-118,120], rounded DOWN toward zero to
    # the nearest 5deg so we stay inside measured travel) so the shoulder can bend the gripper down on
    # either side more aggressively — needed for a lower rest pose. The calibration is a rough guide
    # only (real-convention; sim runs the CAD URDF). wrist_pitch and elbow_pitch keep their original
    # sim ranges: widening wrist_pitch gave the null-space solver enough freedom to DRIFT on the
    # self-check's jog-out-and-back (it returns to a different config), so it stays at +-90.
    "shoulder_pitch": (-115.0, 120.0),
    "shoulder_roll": (-180.0, 180.0),
    "elbow_pitch": (-15.0, 195.0),
    "elbow_roll": (-180.0, 180.0),
    "wrist_pitch": (-90.0, 90.0),
}
GRIPPER_LIMITS = (0.0, 100.0)
# The controller's gripper command is 0 = CLOSED, 100 = OPEN. Some real arms wire the servo the
# other way (normalized 0 = jaw OPEN, 100 = CLOSED — the servo-low=open gotcha). GRIPPER_INVERT sends
# (100 - command) to the motor so "close" closes, and converts measured jaw back on seed. Set per arm
# (calibration gripper_invert / runner --gripper-invert).
GRIPPER_INVERT: bool = False


def joint_limits(name: str) -> tuple[float, float]:
    """(lo, hi) degree limits for an arm joint (see JOINT_LIMITS_DEG)."""
    return JOINT_LIMITS_DEG[name]


# PRIORITISED (null-space) IK over IK_ACTIVE_JOINTS with a numerical Jacobian. The PRIMARY task is
# the gripper TIP position -> the commanded XYZ, solved with the damped pseudo-inverse of the 3x6
# position Jacobian (so a move always drives the TIP to where you command it). The SECONDARY tasks --
# keep the gripper pointing down, and pull the redundant joints back toward the ready posture
# (anti-drift, wrist excluded) -- are projected into the NULL SPACE of the primary, so they shape the
# redundancy WITHOUT ever fighting the position. Consequence (exactly what we want, no special
# casing): wherever the tip is reachable with the gripper down, it stays down; when you reach LONG
# and there's no null space left, the position solve naturally ROTATES THE GRIPPER OUT to extend the
# tip. Single deterministic solve -> no flip-flop/shake, no drift.
IK_NS_ORIENT_GAIN = 2.0  # null-space gain pulling the gripper toward straight-down
IK_NS_POSTURE_GAIN = 0.05  # null-space gain pulling redundant joints toward the ready posture
IK_MAX_ITERS = 200  # full-convergence cap (one-shot solves: self-check, demo, apply_command)
IK_CONTROL_ITERS = 12  # per-tick cap in the real-time control loop (warm-started, so a few suffice)
IK_TOLERANCE = 1e-4
IK_DAMPING = 1e-2  # Levenberg damping on the position pseudo-inverse (stable near singularities)
IK_PROBE_DEG = 0.1  # finite-difference step for the numerical Jacobian
IK_MAX_STEP_DEG = 4.0  # clamp per-iteration joint motion -> smooth, hardware-safe

# Reachability gates: a commanded step is rejected (and reverted) if the solved pose misses
# position by more than MAX_POSITION_RESIDUAL_M, or tilts the gripper more than
# MAX_ORIENTATION_RESIDUAL_DEG off vertical. Position is the hard requirement. The tilt cap bounds
# how far the gripper may rotate out to extend reach: 90 deg = it may go fully horizontal for max
# reach (note that past ~45-60 deg it points too sideways to grasp a piece from above, so the
# extra range is reach-only). Lower this if you want to keep every reachable pose graspable.
MAX_POSITION_RESIDUAL_M = 0.01
MAX_ORIENTATION_RESIDUAL_DEG = 90.0

# Work-surface floor: the target is never driven below this Z (the gripper cannot go
# through the table). The base sits at z=0; 3cm clearance keeps the tip off the surface
# so an operator holding "down" can't press the gripper into the table.
FLOOR_Z_M = 0.03

# The point being controlled is the gripper TIP (the end of the claws), not the wrist. The tip
# sits this far out along the gripper's approach axis (local +Z) from gripper_base_link (measured
# from the claw mesh). Commanding the tip is what lets the IK ROTATE THE GRIPPER OUT when reaching
# long: tilting the approach axis toward horizontal throws the tip up to TOOL_LENGTH_M further out
# than the wrist alone could place it.
TOOL_LENGTH_M = 0.0785

# Gripper command (0..100) -> servo angle (deg). gripper=100 maps to the gripper_servo's full URDF
# travel (its <limit upper>, ~335 deg). This fallback is used only if the URDF carries no limit.
DEFAULT_GRIPPER_OPEN_DEG = 335.0

# Gripper command->servo endpoints (CALIBRATABLE). The mapping is intentionally inverted: command
# 100 = OPEN -> small servo angle (jaw open), command 0 = CLOSED -> large servo angle (jaw shut).
# -m tomctl.calibrate.arm records the real open/closed servo positions here; None = use the URDF travel.
GRIPPER_SERVO_OPEN_DEG = 0.0
GRIPPER_SERVO_CLOSED_DEG: float | None = None

# Whether the gripper is present. The gripper is standard, so this defaults True; an older
# calibration carrying "include_gripper": false can still drop it from the controls/rendering for a
# bare 6-joint arm. The IK is unaffected either way (the gripper was never in the IK chain).
INCLUDE_GRIPPER = True

# A clean, down-facing starting pose (degrees), valid under the calibrated limits. base_yaw=+90
# aims the arm at +X (over the table, away from the green board/clamp on the -X edge); rolls 0; the
# arm stands tall with the wrist pitched down so the gripper points at the floor (~2 deg off
# vertical) and the claw tip hovers ~15.6 cm above the board — the operator-chosen ready height.
READY_POSE_DEG = {
    "base_yaw": 90.0,
    "shoulder_pitch": 0.0,
    "shoulder_roll": 0.0,
    "elbow_pitch": 1.0,
    "elbow_roll": 0.0,
    "wrist_pitch": -89.0,
}

# A stowed REST pose (degrees, operator-chosen via the joint sliders): the shoulder bends well down
# (shoulder_pitch +109, using the widened range) with the elbow folded and forearm rolled so the arm
# lies low and compact to rest. Reached only via HOME (path planner). SIM ONLY — the real arm's rest
# pose comes from ik_calibration.json's rest_pose_deg, which overrides this at load.
RESTING_POSE_DEG = {
    "base_yaw": 7.0,
    "shoulder_pitch": 109.0,
    "shoulder_roll": 76.0,
    "elbow_pitch": 146.0,
    "elbow_roll": -67.0,
    "wrist_pitch": 0.0,
}

# Work surface (the black ~75 cm desk) in world metres. 75 cm is the reach-optimal table length
# for two arms clamped at opposite ends to BOTH reach every square of a 15" board (see the
# board-sizing analysis). The two arms clamp to its two ends and reach in toward the centre. Each
# base is shifted inboard so the clamp's outer wall sits FLUSH with the desk edge (clamp wall is
# ~0.069 m outboard of the base, desk edge at TABLE_CENTER_X -+ 0.375). The desk is thick enough
# that the clamp visibly grips its edge.
ARM1_BASE_X = -0.021
ARM2_BASE_X = 0.591
TABLE_CENTER = (0.285, 0.0, -0.025)  # midway between the arms; top at z=0 (thickness/2 below)
TABLE_DIAMETER = 0.90  # circular table, fixed 90 cm diameter (arms clamp to opposite edges)
TABLE_THICKNESS = 0.05
# Base-origin inboard offset from the clamped table edge: the TOM clamp's outer wall sits flush with
# the edge, ~0.069 m outboard of the base. Used both to place the arms in build_scene and to derive
# the default teleop workspace centre below (so the two stay consistent).
CLAMP_INSET_M = 0.069

# --- Teleop workspace boundary (the table footprint) -------------------------------------
# During MANUAL Cartesian control the gripper TIP's XY is confined to the circular table so a held
# arrow key (especially a base_yaw swing) can NEVER carry the gripper off the table or behind the
# base. The disk lives in the arm's OWN IK frame (base at the origin). NAMED-POSE moves (home / rest
# / custom) interpolate in JOINT space and BYPASS this entirely — they may leave the disk (the stowed
# REST pose sits off to the side), exactly as the operator asked.
#   * WORKSPACE_RADIUS_M  — disk radius (metres); None disables the boundary.
#   * WORKSPACE_CENTER_XY — disk centre in the arm's own frame (metres).
# Defaults match the SIM's 90 cm table as placed by build_scene: each arm is clamped flush to the
# edge and reaches +X, so the table centre is (radius - inset) ≈ 0.381 m straight ahead. The REAL arm
# reaches +Y instead (and its auto jog-yaw is unreliable because HOME is nearly vertical), so it
# OVERRIDES WORKSPACE_CENTER_XY via ik_calibration.json rather than deriving it from the jog frame.
WORKSPACE_RADIUS_M: float | None = TABLE_DIAMETER / 2.0  # 0.45 m (the table radius)
WORKSPACE_CENTER_XY: tuple[float, float] = (TABLE_DIAMETER / 2.0 - CLAMP_INSET_M, 0.0)  # ≈ (0.381, 0)

# CC0 chess set (board + pieces, for show only), vendored under assets/chess/ and served
# locally so its glTF's relative texture paths resolve. Source: https://polyhaven.com/a/chess_set
# (CC0). Centred on the table. CHESS_FIT_SIZE is the OVERALL board footprint (incl. border) the
# model is scaled to. 0.381 m = a 15" board (the target purchase), whose ~34.5 cm playing area /
# ~43 mm squares sits right at the two-arm both-reach limit (~33.5 cm) on the 75 cm table.
CHESS_GLTF_URL = "/assets/chess/chess_set_1k.gltf"
CHESS_FIT_SIZE = 0.381
CHESS_YAW_DEG = 90.0  # spin the board about vertical so the two armies face the two arms (+-X)

# Content types for the vendored chess bundle served under /assets/.
_CONTENT_TYPES = {
    ".gltf": "model/gltf+json",
    ".bin": "application/octet-stream",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".png": "image/png",
}

# Light, neutral background so the grid + dark table read cleanly.
WEB_BG_COLOR = 0xF1F3F5

# Default 3D camera for the web viewer: a 3/4 view from the -X (arm 1) outward side, pulled back
# and raised to frame both arms and the board. up-arrow = +X = away from arm 1 = into the scene.
WEB_CAMERA_EYE = (-0.74, -0.54, 0.70)
WEB_CAMERA_TARGET = (0.285, 0.0, 0.0)

DEFAULT_VISUAL_COLOR = (0.60, 0.63, 0.65, 1.0)  # neutral grey for meshes without a material


# ========================================================================================
# Small transform helpers (same conventions as lerobot.robots.tom.kinematics).
# ========================================================================================
def _rpy_to_matrix(rpy_rad: NDArray[np.float64]) -> NDArray[np.float64]:
    roll, pitch, yaw = rpy_rad
    cr, sr = np.cos(roll), np.sin(roll)
    cp, sp = np.cos(pitch), np.sin(pitch)
    cy, sy = np.cos(yaw), np.sin(yaw)
    rot_x = np.array([[1, 0, 0], [0, cr, -sr], [0, sr, cr]])
    rot_y = np.array([[cp, 0, sp], [0, 1, 0], [-sp, 0, cp]])
    rot_z = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
    return rot_z @ rot_y @ rot_x


def _origin_transform(xyz: NDArray[np.float64], rpy_rad: NDArray[np.float64]) -> NDArray[np.float64]:
    t = np.eye(4)
    t[:3, :3] = _rpy_to_matrix(rpy_rad)
    t[:3, 3] = xyz
    return t


def _axis_rotation(axis: NDArray[np.float64], angle_rad: float) -> NDArray[np.float64]:
    x, y, z = axis / np.linalg.norm(axis)
    c, s = np.cos(angle_rad), np.sin(angle_rad)
    m = 1.0 - c
    t = np.eye(4)
    t[:3, :3] = np.array(
        [
            [c + x * x * m, x * y * m - z * s, x * z * m + y * s],
            [y * x * m + z * s, c + y * y * m, y * z * m - x * s],
            [z * x * m - y * s, z * y * m + x * s, c + z * z * m],
        ]
    )
    return t


def _axis_translation(axis: NDArray[np.float64], distance: float) -> NDArray[np.float64]:
    t = np.eye(4)
    t[:3, 3] = axis * distance
    return t


def _vec(text: str | None, default: tuple[float, ...]) -> NDArray[np.float64]:
    if text is None:
        return np.asarray(default, dtype=np.float64)
    return np.asarray([float(v) for v in text.split()], dtype=np.float64)


# ========================================================================================
# URDF model: links + their visuals, the joint tree, and full-tree forward kinematics.
# ========================================================================================
@dataclass
class Visual:
    mesh_path: Path
    scale: NDArray[np.float64]
    origin: NDArray[np.float64]  # 4x4, link-relative
    color_hex: int


@dataclass
class Joint:
    name: str
    kind: str
    parent: str
    child: str
    origin: NDArray[np.float64]  # 4x4
    axis: NDArray[np.float64]
    mimic: tuple[str, float, float] | None  # (source joint, multiplier, offset)


@dataclass
class UrdfModel:
    urdf_path: Path
    links: dict[str, list[Visual]]
    joints: list[Joint]
    revolute_limits_rad: dict[str, tuple[float, float]] = field(default_factory=dict)
    continuous_joints: tuple[str, ...] = ()  # full-turn joints (base_yaw, shoulder/elbow_roll)
    gripper_servo_max_rad: float | None = None  # gripper_servo <limit upper> = fully open

    @property
    def root(self) -> str:
        children = {j.child for j in self.joints}
        roots = [name for name in self.links if name not in children]
        if len(roots) != 1:
            raise ValueError(f"Expected exactly one root link, found {roots}.")
        return roots[0]

    def children_of(self, link: str) -> list[Joint]:
        return [j for j in self.joints if j.parent == link]

    def forward_kinematics_links(self, joint_values_deg: dict[str, float]) -> dict[str, NDArray[np.float64]]:
        """World transform of every link, given joint values (degrees for revolute/continuous).

        Prismatic mimic joints (the gripper claws) derive their slide from their source
        joint: slide_metres = multiplier * radians(source) + offset.
        """
        frames = {self.root: np.eye(4)}
        stack = [self.root]
        while stack:
            parent = stack.pop()
            for joint in self.children_of(parent):
                frames[joint.child] = frames[parent] @ joint.origin @ self._motion(joint, joint_values_deg)
                stack.append(joint.child)
        return frames

    @staticmethod
    def _motion(joint: Joint, q: dict[str, float]) -> NDArray[np.float64]:
        if joint.kind in ("revolute", "continuous"):
            return _axis_rotation(joint.axis, np.deg2rad(q.get(joint.name, 0.0)))
        if joint.kind == "prismatic":
            if joint.mimic is not None:
                src, mult, off = joint.mimic
                slide = mult * np.deg2rad(q.get(src, 0.0)) + off
            else:
                slide = q.get(joint.name, 0.0)
            return _axis_translation(joint.axis, slide)
        return np.eye(4)  # fixed


def _color_to_hex(rgba: tuple[float, ...]) -> int:
    r, g, b = (int(round(max(0.0, min(1.0, c)) * 255)) for c in rgba[:3])
    return (r << 16) | (g << 8) | b


def load_urdf_model(urdf_path: Path) -> UrdfModel:
    root = ET.parse(urdf_path).getroot()

    materials: dict[str, int] = {}
    for mat in root.findall("material"):
        color = mat.find("color")
        if color is not None and "rgba" in color.attrib:
            materials[mat.attrib["name"]] = _color_to_hex(
                tuple(_vec(color.attrib["rgba"], (0.6, 0.6, 0.6, 1.0)))
            )

    links: dict[str, list[Visual]] = {}
    for link in root.findall("link"):
        visuals: list[Visual] = []
        for vis in link.findall("visual"):
            mesh = vis.find("geometry/mesh")
            if mesh is None:  # only mesh visuals are rendered (Tom is all-mesh)
                continue
            origin = vis.find("origin")
            xyz = _vec(origin.attrib.get("xyz") if origin is not None else None, (0, 0, 0))
            rpy = _vec(origin.attrib.get("rpy") if origin is not None else None, (0, 0, 0))
            scale = _vec(mesh.attrib.get("scale"), (1, 1, 1))
            mat = vis.find("material")
            color = materials.get(mat.attrib["name"]) if mat is not None and "name" in mat.attrib else None
            visuals.append(
                Visual(
                    mesh_path=(urdf_path.parent / mesh.attrib["filename"]).resolve(),
                    scale=scale,
                    origin=_origin_transform(xyz, rpy),
                    color_hex=color if color is not None else _color_to_hex(DEFAULT_VISUAL_COLOR),
                )
            )
        links[link.attrib["name"]] = visuals

    joints: list[Joint] = []
    limits: dict[str, tuple[float, float]] = {}
    gripper_servo_max_rad: float | None = None
    for joint in root.findall("joint"):
        origin = joint.find("origin")
        axis = joint.find("axis")
        mimic_el = joint.find("mimic")
        mimic = None
        if mimic_el is not None:
            mimic = (
                mimic_el.attrib["joint"],
                float(mimic_el.attrib.get("multiplier", 1.0)),
                float(mimic_el.attrib.get("offset", 0.0)),
            )
        name = joint.attrib["name"]
        kind = joint.attrib["type"]
        joints.append(
            Joint(
                name=name,
                kind=kind,
                parent=joint.find("parent").attrib["link"],
                child=joint.find("child").attrib["link"],
                origin=_origin_transform(
                    _vec(origin.attrib.get("xyz") if origin is not None else None, (0, 0, 0)),
                    _vec(origin.attrib.get("rpy") if origin is not None else None, (0, 0, 0)),
                ),
                axis=_vec(axis.attrib.get("xyz") if axis is not None else None, (1, 0, 0)),
                mimic=mimic,
            )
        )
        lim = joint.find("limit")
        if lim is not None and "upper" in lim.attrib:
            if name == "gripper_servo":
                gripper_servo_max_rad = float(lim.attrib["upper"])
            elif kind == "revolute" and "lower" in lim.attrib:
                limits[name] = (float(lim.attrib["lower"]), float(lim.attrib["upper"]))

    return UrdfModel(
        urdf_path=urdf_path,
        links=links,
        joints=joints,
        revolute_limits_rad=limits,
        continuous_joints=tuple(j.name for j in joints if j.kind == "continuous"),
        gripper_servo_max_rad=gripper_servo_max_rad,
    )


# ========================================================================================
# The Cartesian controller: position/gripper commands in, arm joint targets out, gripper down.
# ========================================================================================
def _approach_angle_deg(pose: NDArray[np.float64]) -> float:
    """Angle (deg) between the tool's approach axis (its local +Z) and straight down."""
    return float(np.rad2deg(np.arccos(np.clip(pose[:3, 2] @ DOWN_AXIS, -1.0, 1.0))))


def _tool_point(pose: NDArray[np.float64]) -> NDArray[np.float64]:
    """World position of the gripper TIP (the controlled point): wrist + TOOL_LENGTH along approach."""
    return pose[:3, 3] + TOOL_LENGTH_M * pose[:3, 2]


class CartesianController:
    """Turns XY/Z/gripper commands into joints, keeping the gripper pointing roughly downward.

    IK runs over all six arm joints with damped least squares (numerical Jacobian), clamping
    joint limits each iteration. The error stacks three tasks: position (weighted heavily, so the
    gripper REACHES where you command), approach-down (soft — the gripper stays vertical where it
    can and tips up to MAX_ORIENTATION_RESIDUAL_DEG to reach otherwise), and roll-centering (pulls
    the shoulder/elbow rolls back toward 0). The rolls let the arm rotate onto its side to reach
    points the pure yaw+pitch arm can't; centering them resolves the redundancy so there's no
    null-space drift — they engage only when needed and return to 0. Heading follows the reach.
    """

    def __init__(self, backend, gripper_servo_max_rad: float | None = None):
        # `backend` is the ONE seam that lets the SAME controller drive the browser sim and the real
        # arm: only the forward-kinematics chain differs. It is either an FK backend used directly
        # (anything with `.forward_kinematics(joint_dict)->4x4` + `.joint_names` — a TomKinematics,
        # or a live robot's `.kinematics`), or a sim `UrdfModel` from which we build a TomKinematics
        # over its URDF (back-compat for tomctl.calibrate.arm and the sim's run_* helpers, which pass a model).
        if hasattr(backend, "forward_kinematics") and hasattr(backend, "joint_names"):
            self.kin = backend
        else:
            self.kin = TomKinematics(backend.urdf_path)
        if gripper_servo_max_rad is None:
            gripper_servo_max_rad = getattr(backend, "gripper_servo_max_rad", None)
        self.arm_joints = list(self.kin.joint_names)
        self.gripper_open_deg = (
            np.rad2deg(gripper_servo_max_rad)
            if gripper_servo_max_rad is not None
            else DEFAULT_GRIPPER_OPEN_DEG
        )

        # Posture regularisation target/mask (resolves redundancy, prevents drift): pin ONLY the one
        # genuinely-redundant DOF — the gripper's heading about vertical — by centring the two ROLLS.
        # base_yaw is left free (it must swing to aim the arm at LEFT/RIGHT targets) and the pitches
        # are left free (they must EXTEND to reach forward); constraining either would fight the
        # position, limiting lateral and forward motion. Only the rolls are truly redundant here.
        _posture_joints = ("shoulder_roll", "elbow_roll")
        self._ready_active = np.array(
            [READY_POSE_DEG.get(n, 0.0) for n in IK_ACTIVE_JOINTS], dtype=np.float64
        )
        self._posture_mask = np.array(
            [1.0 if n in _posture_joints else 0.0 for n in IK_ACTIVE_JOINTS], dtype=np.float64
        )

        self.joints_deg = self._ready_joints()
        self.target_pos = _tool_point(self.kin.forward_kinematics(self.joints_deg)).copy()
        self.gripper = 0.0  # 0..100 (closed..open)
        self._solve_to_target()  # settle exactly onto the down-pointing target

        # Intent vs execution: `commanded_*` is what the operator/remote wants (it may jump around or
        # arrive late/jittery); the control loop slews the WORKING target (`target_pos`) toward it at
        # a bounded Cartesian speed/accel and runs the IK each tick, so the motion stays smooth no
        # matter how the commands arrive. See control_step().
        self.commanded_pos = self.target_pos.copy()
        self.commanded_gripper = self.gripper
        self._cart_speed = 0.0  # current Cartesian tip speed (m/s), accel-limited toward target
        self._prev_joint_delta = {}  # zero jog joint-velocity memory (accel limiter)

        # Operator jog frame (see JOG_FRAME_YAW_DEG): rotate the (forward, left) key intent into world
        # XY so UP = out over the table along the arm's HOME heading. Auto = the home tip's azimuth;
        # JOG_FRAME_YAW_DEG overrides. Computed once from the home pose; the live/seeded pose doesn't
        # change it (forward stays "the way the arm faces at home", which is what the operator means).
        if JOG_FRAME_YAW_DEG is not None:
            jog_yaw = np.deg2rad(JOG_FRAME_YAW_DEG)
        else:
            jog_yaw = float(np.arctan2(self.target_pos[1], self.target_pos[0]))
        self._jog_cos, self._jog_sin = float(np.cos(jog_yaw)), float(np.sin(jog_yaw))
        self._jog_left_sign = -1.0 if JOG_FLIP_LR else 1.0  # mirror LEFT for a left-handed world XY
        self.jog_yaw_deg = float(np.rad2deg(jog_yaw))

        # Workspace boundary centre (own-frame XY), snapshotted from the calibration-resolved global so
        # a per-arm table placement applies. The radius is read live (WORKSPACE_RADIUS_M) so it can be
        # toggled off. See _clamp_to_workspace and the WORKSPACE_* config above.
        self._ws_center = np.asarray(WORKSPACE_CENTER_XY, dtype=np.float64)
        # Z-floor, snapshotted per-arm (the two-arm host loads a different floor_z_m per arm before build).
        # --debug lifts it by setting this instance field directly (not the global) so per-arm survives.
        self._floor_z = float(FLOOR_Z_M)

        # Named-pose path planner: when active, _path is a list of joint-dict waypoints the control
        # loop interpolates through (joint space), instead of the manual Cartesian slew. Transitions
        # between named poses always route via HOME (see go_rest/go_custom). _custom is a saved pose.
        self._path: list[dict[str, float]] | None = None
        self._path_idx = 0
        self._path_speed = 0.0  # joint-space cruise speed (deg/s), accel-limited
        # A home/rest move waits here (arm HELD) until the gripper has fully opened, so we RELEASE before
        # moving — never carry/drag an object into home/rest. control_step promotes it to _path once open.
        self._pending_path: list[dict[str, float]] | None = None
        self._pending_elapsed = 0.0  # time spent in the release-before-move wait (for the timeout backstop)
        self._custom: dict[str, float] | None = None

    def _ready_joints(self) -> dict[str, float]:
        return self._postprocess({**dict.fromkeys(self.arm_joints, 0.0), **READY_POSE_DEG})

    # -- 6-DOF IK: tip position + soft approach-down + posture regularisation -------------
    def _full(self, active_vec) -> dict[str, float]:
        """Expand the active-joint vector to a full joint dict."""
        out = dict.fromkeys(self.arm_joints, 0.0)
        for name, value in zip(IK_ACTIVE_JOINTS, active_vec, strict=True):
            out[name] = float(value)
        return out

    def _bounds(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
        # Continuous joints (base_yaw + the two rolls) are free-spinning — give the IK effectively
        # unbounded travel for them so it never clamps a full-turn joint at a phantom +-180 seam.
        lohi = np.array(
            [(-1e9, 1e9) if n in CONTINUOUS_JOINTS else joint_limits(n) for n in IK_ACTIVE_JOINTS],
            dtype=np.float64,
        )
        return lohi[:, 0], lohi[:, 1]

    def _solve_to_target(self, max_iters: int = IK_MAX_ITERS) -> dict[str, float]:
        """Prioritised null-space IK: drive the gripper TIP to self.target_pos (primary), and pursue
        gripper-down + ready-posture only in the primary's null space (so they never fight position).
        max_iters bounds the work: the real-time control loop passes a small value (warm-started, a
        few iterations track the slewed target); one-shot callers use the full convergence cap."""
        lo, hi = self._bounds()
        n = len(IK_ACTIVE_JOINTS)
        q = np.array([self.joints_deg[name] for name in IK_ACTIVE_JOINTS], dtype=np.float64)
        eye = np.eye(n)
        for _ in range(max_iters):
            pose = self.kin.forward_kinematics(self._full(q))
            tip = _tool_point(pose)
            pos_err = self.target_pos - tip  # 3-vector, primary task error (metres)
            ori_err = np.cross(pose[:3, 2], DOWN_AXIS)  # 3-vector -> 0 when gripper points down

            # Numerical Jacobians of tip position (Jp) and approach error (Jo) wrt the active joints.
            jp = np.empty((3, n))
            jo = np.empty((3, n))
            for i in range(n):
                probe = q.copy()
                probe[i] += IK_PROBE_DEG
                p2 = self.kin.forward_kinematics(self._full(probe))
                jp[:, i] = (_tool_point(p2) - tip) / IK_PROBE_DEG
                jo[:, i] = (np.cross(p2[:3, 2], DOWN_AXIS) - ori_err) / IK_PROBE_DEG

            # Primary: damped pseudo-inverse step toward the tip target.
            jp_pinv = jp.T @ np.linalg.inv(jp @ jp.T + (IK_DAMPING**2) * np.eye(3))
            dq = jp_pinv @ pos_err
            # Secondary (projected into the null space of the primary): pull the gripper down and the
            # redundant joints toward the ready posture (wrist excluded so it stays free to rotate out).
            null = eye - jp_pinv @ jp
            secondary = -IK_NS_ORIENT_GAIN * (jo.T @ ori_err)
            secondary += IK_NS_POSTURE_GAIN * self._posture_mask * (self._ready_active - q)
            dq = dq + null @ secondary

            biggest = float(np.max(np.abs(dq)))
            if biggest > IK_MAX_STEP_DEG:
                dq *= IK_MAX_STEP_DEG / biggest
            q = np.clip(q + dq, lo, hi)
            if np.linalg.norm(pos_err) <= IK_TOLERANCE and np.linalg.norm(null @ secondary) <= IK_TOLERANCE:
                break
        self.joints_deg = self._postprocess(self._full(q))
        return self.joints_deg

    def _postprocess(self, joints_deg: dict[str, float]) -> dict[str, float]:
        """Clamp every limit-calibrated arm joint to its range (JOINT_LIMITS_DEG); leave the
        continuous (free-spinning) joints — base_yaw + the two rolls — untouched."""
        out = dict(joints_deg)
        for name in IK_ACTIVE_JOINTS:
            if name in CONTINUOUS_JOINTS:
                continue
            lo, hi = joint_limits(name)
            out[name] = float(np.clip(out[name], lo, hi))
        return out

    # -- residuals of the current solution against the current target --------------------
    def residuals(self) -> tuple[float, float]:
        pose = self.kin.forward_kinematics(self.joints_deg)
        pos_err = float(np.linalg.norm(self.target_pos - _tool_point(pose)))
        return pos_err, _approach_angle_deg(pose)

    # -- public command API --------------------------------------------------------------
    def apply_command(self, dx: float, dy: float, dz: float, dgrip: float) -> bool:
        """Apply a task-space delta to the gripper TIP. Returns False (and reverts) if unreachable.
        The gripper open/close is always applied (independent of reachability)."""
        prev_pos, prev_joints = self.target_pos.copy(), dict(self.joints_deg)

        self.target_pos = self.target_pos + np.array([dx, dy, dz])
        self.target_pos[2] = max(self._floor_z, self.target_pos[2])
        self.gripper = float(np.clip(self.gripper + dgrip, *GRIPPER_LIMITS))

        self._solve_to_target()
        pos_err, approach_err = self.residuals()
        if pos_err > MAX_POSITION_RESIDUAL_M or approach_err > MAX_ORIENTATION_RESIDUAL_DEG:
            # Out of reach / would tilt the gripper past the cap: revert the position change. The
            # gripper open/close is kept — it is independent of reachability.
            self.target_pos, self.joints_deg = prev_pos, prev_joints
            return False
        return True

    # -- workspace boundary (manual control only) ----------------------------------------
    def _clamp_to_workspace(self, xy: NDArray[np.float64]) -> NDArray[np.float64]:
        """Project a tip XY back onto the table disk (centre ``self._ws_center``, radius
        ``WORKSPACE_RADIUS_M``) when it falls outside, so MANUAL jogging can't drive the gripper off
        the table — in particular it stops a base_yaw swing from carrying the tip behind the base / off
        the back edge. Only the horizontal footprint is bounded; Z is untouched. No-op when the
        boundary is disabled (radius None). Named-pose moves never reach this — control_step returns
        early while a joint-space path runs — so home/rest transitions are free to leave the disk."""
        r = WORKSPACE_RADIUS_M
        if r is None:
            return xy
        d = xy - self._ws_center
        dist = float(np.hypot(d[0], d[1]))
        # Defensive: a non-finite tip (NaN/inf from a bad remote target) makes ``dist > r`` False, so
        # the poisoned value would pass straight through into the IK/slew. Collapse it to the (safe,
        # in-bounds) workspace centre instead. The remote receiver also rejects non-finite targets
        # upstream; this is the last-resort guard.
        if not np.isfinite(dist):
            return self._ws_center.copy()
        if dist > r:
            return self._ws_center + d * (r / dist)
        return xy

    # -- intent API (operator/remote) + smooth control step ------------------------------
    def command_delta(self, dx: float, dy: float, dz: float, dgrip: float) -> None:
        """Advance the COMMANDED (intent) tip target + gripper. dx/dy are in the OPERATOR jog frame
        (dx = forward/UP, dy = left/LEFT); they're rotated into world XY by the jog-frame yaw so UP
        moves out over the table along the arm's home heading. Does NOT move the arm directly — the
        control loop slews the working target toward it and runs the IK (see control_step)."""
        self._cancel_path()  # manual input takes over from any running pose-to-pose move
        left = self._jog_left_sign * dy  # mirror LEFT when the world XY is left-handed (JOG_FLIP_LR)
        wx = dx * self._jog_cos - left * self._jog_sin  # operator (forward, left) -> world (x, y)
        wy = dx * self._jog_sin + left * self._jog_cos
        self.commanded_pos = self.commanded_pos + np.array([wx, wy, dz])
        self.commanded_pos[:2] = self._clamp_to_workspace(self.commanded_pos[:2])  # keep on the table
        self.commanded_pos[2] = max(self._floor_z, self.commanded_pos[2])
        if INCLUDE_GRIPPER:  # no-op the gripper channel when there's no gripper
            self.commanded_gripper = float(np.clip(self.commanded_gripper + dgrip, *GRIPPER_LIMITS))

    # -- named-pose path planning (home / rest / custom; always via home) ----------------
    def _pose(self, pose_deg: dict[str, float]) -> dict[str, float]:
        return self._postprocess({**dict.fromkeys(self.arm_joints, 0.0), **pose_deg})

    def _start_path(self, waypoints: list[dict[str, float]]) -> None:
        self._path = waypoints
        self._path_idx = 0
        self._path_speed = 0.0

    def _cancel_path(self) -> None:
        self._path = None

    def _raised(self, base_yaw: float) -> dict[str, float]:
        """The HOME (raised) arm configuration but with base_yaw set to a given heading. Used to
        sequence rest<->home so base_yaw only ever turns while the arm is UP (clear of the table)."""
        return self._pose({**READY_POSE_DEG, "base_yaw": float(base_yaw)})

    def go_home(self) -> None:
        """Smoothly move to the HOME (ready) pose in TWO stages: first RAISE to the home configuration
        with base_yaw held fixed at its current heading (a low/rest pose lifts straight up, no
        horizontal sweep), THEN rotate base_yaw to the home heading while the arm is up.

        RELEASES the gripper FIRST (opens fully), THEN moves — control_step holds the arm until the
        gripper is open, so a HOME never carries/drags an object."""
        self.commanded_gripper = GRIPPER_LIMITS[1]  # 100 = open = release anything held
        self._path = None  # cancel any current move; hold until released (see control_step)
        self._pending_elapsed = 0.0
        self._pending_path = [self._raised(self.joints_deg["base_yaw"]), self._pose(READY_POSE_DEG)]

    def go_rest(self, via_home: bool = True) -> None:
        """Move to the stowed REST pose. via_home (the operator/park path) goes in THREE stages so
        base_yaw only turns while UP: RAISE to the home config with base_yaw held fixed, then ROTATE
        base_yaw to the rest heading (still up), then LOWER into rest. via_home=False goes DIRECTLY to
        rest — used at start-up, where the arm is already at/near rest, to avoid a lift-and-lower.

        RELEASES the gripper FIRST (opens fully), THEN moves — so we never stow holding an object."""
        self.commanded_gripper = GRIPPER_LIMITS[1]  # 100 = open = release anything held
        self._path = None  # cancel any current move; hold until released (see control_step)
        self._pending_elapsed = 0.0
        rest = self._pose(RESTING_POSE_DEG)
        if not via_home:
            self._pending_path = [rest]
            return
        self._pending_path = [
            self._raised(self.joints_deg["base_yaw"]), self._raised(RESTING_POSE_DEG["base_yaw"]), rest
        ]

    def save_custom(self) -> None:
        """Remember the current pose as the CUSTOM pose (recallable later via go_custom)."""
        self._custom = dict(self.joints_deg)

    def go_custom(self) -> None:
        """Smoothly move to the saved CUSTOM pose, ALWAYS via home. No-op if none saved."""
        if self._custom is not None:
            self._start_path([self._pose(READY_POSE_DEG), dict(self._custom)])

    def is_settling(self) -> bool:
        """True while a home/rest sequence is still running — the gripper-release WAIT or the joint move.
        Lets a caller wait for the full park (release + move) to finish before releasing torque."""
        return self._path is not None or self._pending_path is not None

    def has_custom(self) -> bool:
        return self._custom is not None

    def _step_path(self, dt: float) -> None:
        """Advance one tick of a planned joint-space move: trapezoidal speed, all joints synchronised
        so they arrive together; on reaching the final waypoint, hand back to manual Cartesian control."""
        target = self._path[self._path_idx]
        # CONTINUOUS joints (base_yaw + the two rolls) are free-spinning: interpolate them the SHORT
        # way around (wrap the delta to [-180,180]). Without this a roll near the +-180 seam sweeps the
        # LONG way (up to ~360deg) toward the target — which over-rotated the cabled rolls and slipped
        # past the shortest-angle roll guard. Limit-calibrated joints use the plain delta.
        deltas = {}
        for n in self.arm_joints:
            d = target[n] - self.joints_deg[n]
            if n in CONTINUOUS_JOINTS:
                d = (d + 180.0) % 360.0 - 180.0
            deltas[n] = d
        far = max(abs(d) for d in deltas.values())
        if far < 1e-3:  # waypoint reached
            self.joints_deg = {n: float(target[n]) for n in self.arm_joints}
            self._path_idx += 1
            self._path_speed = 0.0
            if self._path_idx >= len(self._path):  # whole path done -> resume manual from here
                self.target_pos = _tool_point(self.kin.forward_kinematics(self.joints_deg)).copy()
                self.commanded_pos = self.target_pos.copy()
                self._cart_speed = 0.0
                self._prev_joint_delta = {}  # zero jog joint-velocity memory (accel limiter)
                self._path = None
            return
        # trapezoidal profile: cruise at JOINT_PATH_SPEED, but decelerate to stop at the waypoint.
        decel_cap = float(np.sqrt(2.0 * JOINT_PATH_ACCEL_DPS2 * far))
        desired = min(JOINT_PATH_SPEED_DPS, decel_cap)
        self._path_speed += float(
            np.clip(desired - self._path_speed, -JOINT_PATH_ACCEL_DPS2 * dt, JOINT_PATH_ACCEL_DPS2 * dt)
        )
        self._path_speed = max(0.0, self._path_speed)
        move = min(self._path_speed * dt, far)  # degrees the slowest (farthest) joint moves this tick
        s = move / far
        self.joints_deg = {n: self.joints_deg[n] + deltas[n] * s for n in self.arm_joints}

    def command_absolute(self, x: float, y: float, z: float, gripper: float) -> None:
        """Set the COMMANDED tip target + gripper absolutely (used by the remote/UDP path)."""
        self.commanded_pos = np.array([x, y, max(self._floor_z, z)], dtype=np.float64)
        self.commanded_pos[:2] = self._clamp_to_workspace(self.commanded_pos[:2])  # keep on the table
        self.commanded_gripper = float(np.clip(gripper, *GRIPPER_LIMITS))

    def control_step(self, dt: float) -> None:
        """One control tick: slew the working target toward the commanded target with bounded
        speed/accel, solve the IK, and rate-limit the joints — so motion stays smooth however the
        commands arrive (bursty keyboard, or delayed/jittery network targets)."""
        if self._pending_path is not None:
            # RELEASE-BEFORE-MOVE: hold the arm and open the gripper, then start the queued move — so a
            # home/rest never carries or drags a held object. Slew toward a FIXED open target
            # (GRIPPER_LIMITS[1]), NOT the operator's live commanded_gripper: a held grip-CLOSE jog during
            # the transition would otherwise drag the target below the release threshold, the move would
            # never start, and the arm would sit wedged (with is_settling() stuck True, which also wedges
            # the one-arm-at-a-time cross-arm gate). Promote after a bounded timeout regardless, so a
            # stuck/miscalibrated gripper can never block the move indefinitely.
            self._pending_elapsed += dt
            open_to = GRIPPER_LIMITS[1]
            g_to = open_to - self.gripper
            self.gripper += float(np.clip(g_to, -GRIPPER_SPEED_PPS * dt, GRIPPER_SPEED_PPS * dt))
            self.gripper = float(np.clip(self.gripper, *GRIPPER_LIMITS))
            released = self.gripper >= open_to - 0.5
            if released or self._pending_elapsed >= PENDING_RELEASE_TIMEOUT_S:
                if not released:  # backstop fired: log it — a normal release completes in <0.5s
                    print(f"[ctrl] gripper release timed out at {self._pending_elapsed:.1f}s "
                          f"(grip={self.gripper:.0f}/{open_to:.0f}); starting move anyway")
                self._start_path(self._pending_path)
                self._pending_path = None
            return
        if self._path is not None:  # a named-pose move is in progress (joint-space, via home)
            self._step_path(dt)
            return
        # Gripper: slew toward the commanded value at a bounded rate, then keep the EXECUTED value
        # within the calibrated travel (the seed widens GRIPPER_LIMITS to admit a measured jaw, so
        # this clamp never snaps a real position — it only bounds the commanded/sent value).
        g_to = self.commanded_gripper - self.gripper
        self.gripper += float(np.clip(g_to, -GRIPPER_SPEED_PPS * dt, GRIPPER_SPEED_PPS * dt))
        self.gripper = float(np.clip(self.gripper, *GRIPPER_LIMITS))

        # Cartesian: accel/decel-limited slew of the working target toward the commanded target.
        to_go = self.commanded_pos - self.target_pos
        dist = float(np.linalg.norm(to_go))
        if dist <= 1e-4:  # essentially arrived (<0.1 mm): HOLD — don't re-solve the redundant IK each
            self.target_pos = self.commanded_pos.copy()  # tick, which lets base_yaw wander -> jitter
            self._cart_speed = 0.0
            self._prev_joint_delta = {}  # zero jog joint-velocity memory (accel limiter)
            return
        desired = min(CART_MAX_SPEED_MPS, dist / dt)  # ease off near the target (no overshoot)
        self._cart_speed += float(
            np.clip(desired - self._cart_speed, -CART_MAX_ACCEL_MPS2 * dt, CART_MAX_ACCEL_MPS2 * dt)
        )
        self._cart_speed = max(0.0, self._cart_speed)
        step = min(self._cart_speed * dt, dist)
        new_target = self.target_pos + to_go * (step / dist)
        new_target[:2] = self._clamp_to_workspace(new_target[:2])  # authoritative: never send off-table
        new_target[2] = max(self._floor_z, new_target[2])

        # Solve the IK to the slewed target, then handle the two failure modes differently:
        #  * gripper TILTING past its cap -> HARD STOP (hold where we are): flipping the gripper over is
        #    never what the operator wants, so freezing is the safe response.
        #  * position OUT OF REACH -> DON'T freeze. The IK already drove the gripper to the NEAREST
        #    reachable point, so KEEP that best-effort pose and re-anchor the working target onto the tip
        #    it actually reached. Continuing to press then SLIDES the gripper along the reachable
        #    boundary (an arc around the base) instead of stalling — the 3-D analogue of the table-edge
        #    slide. Re-anchoring the commanded intent too (below) bounds it so a held key can't run the
        #    target off to infinity. Mirrors how _clamp_to_workspace projects onto the table rim.
        prev_joints, prev_target = dict(self.joints_deg), self.target_pos.copy()
        self.target_pos = new_target
        self._solve_to_target(max_iters=IK_CONTROL_ITERS)
        pos_err, approach_err = self.residuals()
        if approach_err > MAX_ORIENTATION_RESIDUAL_DEG:  # would tilt the gripper past its cap -> hold
            self.target_pos, self.joints_deg = prev_target, prev_joints
            self.commanded_pos = self.target_pos.copy()
            self._cart_speed = 0.0
            self._prev_joint_delta = {}  # zero jog joint-velocity memory (accel limiter)
            return
        reach_blocked = pos_err > MAX_POSITION_RESIDUAL_M
        if reach_blocked:  # re-anchor the target onto the nearest-reachable tip the IK reached
            new_target = _tool_point(self.kin.forward_kinematics(self.joints_deg)).copy()
            new_target[:2] = self._clamp_to_workspace(new_target[:2])
            new_target[2] = max(self._floor_z, new_target[2])
            self.target_pos = new_target

        # Per-joint speed backstop: if any joint would move faster than JOINT_MAX_SPEED_DPS this tick,
        # scale the whole step down uniformly (keeps the pose consistent; next ticks catch up).
        cap = JOINT_MAX_SPEED_DPS * dt
        biggest = max(abs(self.joints_deg[n] - prev_joints[n]) for n in self.arm_joints)
        if biggest > cap:
            s = cap / biggest
            self.joints_deg = {
                n: prev_joints[n] + (self.joints_deg[n] - prev_joints[n]) * s for n in self.joints_deg
            }
            self.target_pos = prev_target + (new_target - prev_target) * s
            self._cart_speed *= s

        # Per-joint ACCELERATION backstop: cap how much each joint's per-tick step may CHANGE from the
        # last tick, so the IK's redundancy/null-space corrections can't jerk a joint (this is what the
        # manual jog lacked vs the smooth planned-path move). Uniform scale keeps the solved pose; the
        # joint continues last tick's velocity plus a bounded change, then re-anchor the target onto the
        # tip we actually reached so the slew keeps flowing and 'arrived' isn't declared while joints lag.
        accel_cap = JOINT_MAX_ACCEL_DPS2 * dt * dt
        dd = {n: (self.joints_deg[n] - prev_joints[n]) - self._prev_joint_delta.get(n, 0.0) for n in self.joints_deg}
        worst_dd = max((abs(v) for v in dd.values()), default=0.0)
        if worst_dd > accel_cap > 0.0:
            sa = accel_cap / worst_dd
            self.joints_deg = {
                n: prev_joints[n] + self._prev_joint_delta.get(n, 0.0) + dd[n] * sa for n in self.joints_deg
            }
            self.target_pos = _tool_point(self.kin.forward_kinematics(self.joints_deg)).copy()
        self._prev_joint_delta = {n: self.joints_deg[n] - prev_joints[n] for n in self.joints_deg}

        if reach_blocked:
            # Bound the intent to the reachable tip (anti-runaway, like the table clamp): a held key
            # can't push the commanded target past the boundary. The next press re-adds its delta and
            # the IK re-projects, so the gripper keeps sliding along the boundary. Note we do NOT zero
            # _cart_speed here (unlike the orientation hold) — that would re-ramp from rest every tick
            # and make the slide crawl; the small per-tick to_go keeps the motion bounded anyway.
            self.commanded_pos = self.target_pos.copy()

    def reset(self) -> None:
        self.joints_deg = self._ready_joints()
        self.target_pos = _tool_point(self.kin.forward_kinematics(self.joints_deg)).copy()
        self.gripper = 0.0
        self._solve_to_target()
        self.commanded_pos = self.target_pos.copy()
        self.commanded_gripper = self.gripper
        self._cart_speed = 0.0
        self._prev_joint_delta = {}  # zero jog joint-velocity memory (accel limiter)
        self._path = None

    def set_joint(self, name: str, value: float) -> None:
        """Set one joint (or the gripper) directly, clamped to its limit — bypasses IK.

        Used by the joint sliders to verify limits. After moving a joint we resync target_pos to
        the new end-effector position so a subsequent Cartesian command continues smoothly from
        the manually-set pose (no IK jump back to a stale target).
        """
        if name == "gripper":
            self.gripper = float(np.clip(value, *GRIPPER_LIMITS))
            return
        if name not in JOINT_LIMITS_DEG:
            return
        lo, hi = joint_limits(name)
        self.joints_deg[name] = float(np.clip(value, lo, hi))
        self.target_pos = _tool_point(self.kin.forward_kinematics(self.joints_deg)).copy()
        self.commanded_pos = self.target_pos.copy()  # control loop holds here, doesn't slew away
        self._cart_speed = 0.0
        self._prev_joint_delta = {}  # zero jog joint-velocity memory (accel limiter)

    def joint_state(self) -> dict[str, float]:
        """Current arm joint angles + gripper (for the sliders / readouts)."""
        state = {n: float(self.joints_deg[n]) for n in IK_ACTIVE_JOINTS}
        state["gripper"] = float(self.gripper)
        return state

    def gripper_servo_deg(self) -> float:
        # gripper command 0..100 = CLOSED..OPEN, linearly mapped to the calibrated servo endpoints.
        # The mapping is inverted (servo 0 = jaw OPEN ~86 mm, servo max = jaw CLOSED ~11 mm — verified
        # from the claw meshes); -m tomctl.calibrate.arm can override the endpoints for the real gripper.
        closed = self.gripper_open_deg if GRIPPER_SERVO_CLOSED_DEG is None else GRIPPER_SERVO_CLOSED_DEG
        return closed + (self.gripper / 100.0) * (GRIPPER_SERVO_OPEN_DEG - closed)

    def render_joint_values(self) -> dict[str, float]:
        """Full joint map (arm + gripper servo) for the renderer's forward kinematics."""
        values = dict(self.joints_deg)
        values["gripper_servo"] = self.gripper_servo_deg()
        return values

    def action(self) -> dict[str, float]:
        """The `<joint>.pos` action dict you would send the real follower. The gripper command (0=closed,
        100=open) is inverted to the servo's convention when GRIPPER_INVERT (servo-low = jaw open)."""
        act = {f"{name}.pos": self.joints_deg[name] for name in self.arm_joints}
        act["gripper_servo.pos"] = (100.0 - self.gripper) if GRIPPER_INVERT else self.gripper
        return act


# ========================================================================================
# meshcat rendering.
# ========================================================================================
class MeshcatArm:
    """Loads the arm meshes once and updates per-link transforms each frame."""

    def __init__(self, model: UrdfModel, open_browser: bool = True):
        try:
            import meshcat
            import meshcat.geometry as g
        except ImportError as e:
            raise SystemExit(
                "The live sim needs meshcat. Install it, e.g. `uv run --with meshcat python ...` "
                "(it also ships with the `all` / `unitree_g1` extras). The headless `--check` "
                "mode needs no meshcat."
            ) from e

        self.model = model
        self.g = g
        self.vis = meshcat.Visualizer()
        self.root_path = "tom"

        # Static work surface (the circular table; top at z=0). meshcat has no disk primitive, so the
        # box is just a stand-in here; the web viewer renders the real cylinder.
        self.vis["surface"].set_object(
            g.Box([TABLE_DIAMETER, TABLE_DIAMETER, TABLE_THICKNESS]),
            g.MeshLambertMaterial(color=0x14161B, opacity=0.85, transparent=True),
        )
        surface = np.eye(4)
        surface[:3, 3] = list(TABLE_CENTER)
        self.vis["surface"].set_transform(surface)

        # Target marker (where the operator is commanding the gripper to).
        self.vis["target"].set_object(
            g.Sphere(0.008), g.MeshLambertMaterial(color=0x37C871, opacity=0.6, transparent=True)
        )

        # Load every visual once; bake (visual-origin @ mesh-scale) into the visual node.
        # Only links that actually get a node are updated later (see `update`).
        self.rendered_links: list[str] = []
        for link, visuals in model.links.items():
            if not visuals:
                continue
            for i, vis in enumerate(visuals):
                try:
                    geom = g.StlMeshGeometry.from_file(str(vis.mesh_path))
                except Exception as e:
                    raise SystemExit(
                        f"Failed to load mesh {vis.mesh_path} for link '{link}'. "
                        "Rebuild the bundle with build_sim_assets.py.\n  "
                        f"{type(e).__name__}: {e}"
                    ) from e
                node = self.vis[self.root_path][link][f"v{i}"]
                node.set_object(geom, g.MeshLambertMaterial(color=vis.color_hex))
                scale = np.diag([*vis.scale, 1.0])
                node.set_transform(vis.origin @ scale)
            self.rendered_links.append(link)

        if open_browser:
            self.vis.open()

    def url(self) -> str:
        return self.vis.url()

    def update(self, joint_values_deg: dict[str, float], target_pos: NDArray[np.float64]) -> None:
        frames = self.model.forward_kinematics_links(joint_values_deg)
        for link in self.rendered_links:
            if link in frames:
                self.vis[self.root_path][link].set_transform(frames[link])
        marker = np.eye(4)
        marker[:3, 3] = target_pos
        self.vis["target"].set_transform(marker)


# ========================================================================================
# Headless self-check: script each control, verify down-orientation + reachability hold.
# ========================================================================================
def run_check(model: UrdfModel) -> bool:
    ctrl = CartesianController(model)
    print("=== Tom Cartesian sim — headless self-check ===")
    print(f"URDF: {model.urdf_path}")
    print(f"Arm joints: {ctrl.arm_joints}")

    # FK self-consistency: the renderer's all-link FK must agree with the IK helper's
    # FK at the target frame (both read the same URDF).
    frames = model.forward_kinematics_links(ctrl.render_joint_values())
    ee_render = frames["gripper_base_link"]
    ee_kin = ctrl.kin.forward_kinematics(ctrl.joints_deg)
    fk_match = float(np.max(np.abs(ee_render - ee_kin)))
    print(
        f"\n[FK ] renderer vs IK-helper EE pose max abs diff: {fk_match:.2e}  "
        f"({'OK' if fk_match < 1e-6 else 'MISMATCH'})"
    )

    start = ctrl.kin.forward_kinematics(ctrl.joints_deg)[:3, 3]
    print(f"[EE ] ready position (m): {np.round(start, 4).tolist()}")
    pos0, ori0 = ctrl.residuals()
    print(f"[EE ] ready residuals: pos={pos0 * 1000:.2f} mm, orient={ori0:.2f} deg")

    # Script each command group; track worst-case off-vertical tilt and any rejections.
    sequences = [
        ("XY +x (right)", {"dx": STEP_XY_M}, 12),
        ("XY -x (left)", {"dx": -STEP_XY_M}, 24),
        ("XY -y (away)", {"dy": -STEP_XY_M}, 12),
        ("XY +y (toward)", {"dy": STEP_XY_M}, 12),
        ("Z down", {"dz": -STEP_Z_M}, 8),
        ("Z up", {"dz": STEP_Z_M}, 8),
        ("gripper open", {"dgrip": STEP_GRIPPER}, 20),
        ("gripper close", {"dgrip": -STEP_GRIPPER}, 20),
    ]
    ok = fk_match < 1e-6
    worst_ori = ori0
    worst_violation = 0.0  # worst amount any arm joint went outside its allowed limits
    print("\n[RUN] scripting each control group:")
    for label, delta, n in sequences:
        rejected = 0
        seq_ori = 0.0
        full = {"dx": 0.0, "dy": 0.0, "dz": 0.0, "dgrip": 0.0}
        full.update(delta)
        for _ in range(n):
            if not ctrl.apply_command(**full):
                rejected += 1
            _, ori = ctrl.residuals()
            seq_ori = max(seq_ori, ori)
            worst_ori = max(worst_ori, ori)
            # Only the limit-calibrated joints have a range to violate; the continuous joints
            # (base_yaw + the two rolls) are deliberately unclamped, so skip them here too.
            for j in LIMIT_CALIBRATED_JOINTS:
                lo, hi = joint_limits(j)
                worst_violation = max(worst_violation, lo - ctrl.joints_deg[j], ctrl.joints_deg[j] - hi)
        ee = ctrl.kin.forward_kinematics(ctrl.joints_deg)[:3, 3]
        flag = "ok" if rejected == 0 else f"{rejected}/{n} rejected (reached workspace edge)"
        print(
            f"   {label:16} -> EE={np.round(ee, 3).tolist()}  grip={ctrl.gripper:5.1f}  "
            f"tilt={seq_ori:.2f} deg  [{flag}]"
        )

    held = worst_ori <= MAX_ORIENTATION_RESIDUAL_DEG
    ok = ok and held

    # Drift check: jog far out and back to the home target; the arm must return to (near) its
    # start pose. With the rolls regularised toward 0 this stays stable — the regression guard
    # for the old null-space drift. (max joint difference over all arm joints after out-and-back)
    ctrl.reset()
    home = dict(ctrl.joints_deg)
    for _ in range(30):
        ctrl.apply_command(STEP_XY_M, 0.0, 0.0, 0.0)
    for _ in range(30):
        ctrl.apply_command(-STEP_XY_M, 0.0, 0.0, 0.0)
    drift = max(abs(ctrl.joints_deg[j] - home[j]) for j in ctrl.arm_joints)
    drift_ok = drift < 2.0
    print(
        f"\n[DRIFT] after a +X/-X excursion: max joint return error {drift:.2f} deg "
        f"({'OK' if drift_ok else 'DRIFTED'})"
    )
    ok = ok and drift_ok

    # Every arm joint must stay within its calibrated range (the IK clamps to joint_limits()).
    limits_ok = worst_violation <= 1e-6
    print(
        f"[LIMITS] all joints within their ranges during run "
        f"(worst overshoot {worst_violation:.3f} deg) -> {'OK' if limits_ok else 'EXCEEDED'}"
    )
    ok = ok and limits_ok

    # Gripper claws must move with the command. Command 0..100 = CLOSED..OPEN; the servo mapping
    # inverts internally (servo 0 = jaw OPEN ~86 mm, servo max = jaw CLOSED ~11 mm), so the OPEN
    # command must drive the servo toward 0.
    ctrl.gripper = 100.0
    open_servo = ctrl.gripper_servo_deg()
    ctrl.gripper = 0.0
    closed_servo = ctrl.gripper_servo_deg()
    open_links = model.forward_kinematics_links({**READY_POSE_DEG, "gripper_servo": open_servo})
    closed_links = model.forward_kinematics_links({**READY_POSE_DEG, "gripper_servo": closed_servo})
    claw_travel = float(
        np.linalg.norm(
            open_links["gripper_right_claw_link"][:3, 3] - closed_links["gripper_right_claw_link"][:3, 3]
        )
    )
    grip_ok = claw_travel > 1e-3 and open_servo < closed_servo  # OPEN command -> servo toward 0 (jaw opens)
    print(
        f"\n[GRIP] claws move {claw_travel * 1000:.1f} mm; OPEN cmd -> servo {open_servo:.0f} deg, "
        f"CLOSE cmd -> servo {closed_servo:.0f} deg  ({'OK' if grip_ok else 'WRONG DIRECTION'})"
    )
    ok = ok and grip_ok

    print(
        f"[ORI  ] worst-case gripper tilt off vertical: {worst_ori:.2f} deg "
        f"(limit {MAX_ORIENTATION_RESIDUAL_DEG:.0f}) -> {'stayed down' if held else 'TILTED'}"
    )
    print(f"\n=== {'PASS' if ok else 'REVIEW — see flags above'} ===")
    return ok


# ========================================================================================
# Keyboard bindings. Each binding maps a canonical key name (from lerobot.utils.keyboard_input:
# arrows -> 'up'/'down'/'left'/'right', else the character) to a (field, delta) command-state
# edit. Deltas are in each arm's OWN frame: +dx = out over the table (away from that arm's base),
# +dy = that arm's left. Each function has its OWN single key (no shared/combo keys):
#   Arm 1 (right hand) = arrows (XY) + ' / . (height up/down) + [ / ] (grip open/close) + H (home)
#   Arm 2 (left hand)  = WASD  (XY) + Q / E (height up/down) + R / F (grip open/close) + B (home)
# ========================================================================================
ARM1_BIND = {
    "up": ("dx", +STEP_XY_M),
    "down": ("dx", -STEP_XY_M),
    "left": ("dy", +STEP_XY_M),
    "right": ("dy", -STEP_XY_M),
    "'": ("dz", +STEP_Z_M),
    ".": ("dz", -STEP_Z_M),
    "[": ("dgrip", +STEP_GRIPPER),
    "]": ("dgrip", -STEP_GRIPPER),
    "h": ("home", None),  # smooth path to home
    # named-pose buttons (browser sends these multi-char codes; keyboard sends 'h')
    "rest1": ("rest", None),
    "save1": ("save", None),
    "cust1": ("custom", None),
}
ARM2_BIND = {
    "w": ("dx", +STEP_XY_M),
    "s": ("dx", -STEP_XY_M),
    "a": ("dy", +STEP_XY_M),
    "d": ("dy", -STEP_XY_M),
    "q": ("dz", +STEP_Z_M),
    "e": ("dz", -STEP_Z_M),
    "r": ("dgrip", +STEP_GRIPPER),
    "f": ("dgrip", -STEP_GRIPPER),
    "b": ("home", None),  # smooth path to home
    "rest2": ("rest", None),
    "save2": ("save", None),
    "cust2": ("custom", None),
}


def new_cmd_state() -> dict:
    return {"dx": 0.0, "dy": 0.0, "dz": 0.0, "dgrip": 0.0, "reset": False, "quit": False}


def _apply_binding(key: str, state: dict, bind: dict) -> bool:
    if key not in bind:
        return False
    field, delta = bind[key]
    if field == "reset":
        state["reset"] = True
    else:
        state[field] += delta
    return True


def _apply_key_intent(key: str, ctrl: "CartesianController", bind: dict) -> bool:
    """Route a key to a controller's INTENT (commanded target), per a binding. Each press nudges the
    commanded tip target / gripper; the control loop slews toward it smoothly. Returns False if the
    key isn't in this arm's binding (so the caller can try the other arm)."""
    if key not in bind:
        return False
    field, delta = bind[key]
    if field == "home":
        ctrl.go_home()
    elif field == "rest":
        ctrl.go_rest()
    elif field == "save":
        ctrl.save_custom()
    elif field == "custom":
        ctrl.go_custom()
    elif field == "reset":
        ctrl.reset()
    elif field == "dx":
        ctrl.command_delta(delta, 0.0, 0.0, 0.0)
    elif field == "dy":
        ctrl.command_delta(0.0, delta, 0.0, 0.0)
    elif field == "dz":
        ctrl.command_delta(0.0, 0.0, delta, 0.0)
    elif field == "dgrip":
        ctrl.command_delta(0.0, 0.0, 0.0, delta)
    return True


def apply_key(name: str, state: dict) -> None:
    """Accumulate one key press into a single-arm command ``state`` (arm-1 bindings).

    Kept module-level and pure so it is unit-testable without a real keyboard backend.
    Used by the single-arm meshcat modes; the web mode routes keys to two arms (see run_web).
    """
    key = name.lower()
    if key in ("q", "esc"):
        state["quit"] = True
    else:
        _apply_binding(key, state, ARM1_BIND)


def run_live(model: UrdfModel) -> None:
    from lerobot.utils.keyboard_input import TerminalKeyListener, create_key_listener

    ctrl = CartesianController(model)
    arm = MeshcatArm(model, open_browser=True)
    arm.update(ctrl.render_joint_values(), ctrl.target_pos)
    print("\nMeshcat viewer:", arm.url())

    state = {"dx": 0.0, "dy": 0.0, "dz": 0.0, "dgrip": 0.0, "reset": False, "quit": False}
    lock = threading.Lock()

    def dispatch(name: str) -> None:
        with lock:
            apply_key(name, state)

    listener = create_key_listener(dispatch, controls_help="arrows/z/x/o/c/h/q")
    if listener is None:
        print("No keyboard backend available (headless / non-interactive). Try --demo or --check.")
        return

    controls = (
        "  arrows    move gripper in the XY surface plane (fixed height)\n"
        "  ' / .     move gripper UP / DOWN (world Z)\n"
        "  [ / ]     open / close gripper\n"
        "  h         home pose        q / Esc    quit"
    )
    # Tell the user which backend won and, crucially, WHERE keypresses are read — the #1
    # confusion is pressing keys with the browser focused while a terminal listener is active.
    if isinstance(listener, TerminalKeyListener):
        print(
            "\nKeyboard: TERMINAL capture — keep THIS TERMINAL window focused (NOT the browser).\n"
            "  To drive it while watching the browser, grant your terminal app 'Input Monitoring'\n"
            "  under System Settings > Privacy & Security, then rerun (enables global capture)."
        )
    else:
        print("\nKeyboard: GLOBAL capture (pynput) — keys work with the BROWSER focused.")
    print(f"\nControls:\n{controls}\n")

    dt = 1.0 / LOOP_HZ
    frame = 0
    try:
        while True:
            with lock:
                cmd = (state["dx"], state["dy"], state["dz"], state["dgrip"])
                do_reset, do_quit = state["reset"], state["quit"]
                state.update(dx=0.0, dy=0.0, dz=0.0, dgrip=0.0, reset=False)
            if do_quit:
                break
            if do_reset:
                ctrl.reset()
            elif any(cmd):
                ctrl.apply_command(*cmd)
            arm.update(ctrl.render_joint_values(), ctrl.target_pos)
            frame += 1
            if frame % 6 == 0:  # ~5 Hz live status so keypresses visibly register
                p = ctrl.target_pos
                print(
                    f"\r  EE x={p[0]:+.3f} y={p[1]:+.3f} z={p[2]:+.3f} m   grip={ctrl.gripper:5.1f}%    ",
                    end="",
                    flush=True,
                )
            time.sleep(dt)
    except KeyboardInterrupt:
        pass
    finally:
        # Ignore a second Ctrl-C while restoring the terminal, so the TTY can't be left in
        # cbreak mode (which would need `stty sane`). keyboard_input also registers an atexit
        # restore as a backstop.
        prev_sigint = signal.signal(signal.SIGINT, signal.SIG_IGN)
        try:
            listener.stop()
        finally:
            signal.signal(signal.SIGINT, prev_sigint)
        print("\nStopped.")


# ========================================================================================
# Scripted demo: same controller, no keyboard. Plays a representative trajectory on a loop
# so you can watch the arm obey every control (XY / Z / gripper / yaw) with the gripper
# held down. Useful for a hands-off look or where keyboard capture isn't available.
# ========================================================================================
def run_demo(model: UrdfModel) -> None:
    ctrl = CartesianController(model)
    arm = MeshcatArm(model, open_browser=True)
    arm.update(ctrl.render_joint_values(), ctrl.target_pos)
    print("\nMeshcat viewer:", arm.url())
    print("Playing scripted demo (gripper stays down throughout). Ctrl-C to stop.\n")

    # (label, dx, dy, dz, dgrip, steps). +X is "out over the table", +-Y is sideways.
    script = [
        ("reach out (+X)", STEP_XY_M, 0.0, 0.0, 0.0, 16),
        ("lower to surface", 0.0, 0.0, -STEP_Z_M, 0.0, 12),
        ("open gripper", 0.0, 0.0, 0.0, STEP_GRIPPER, 20),
        ("close gripper", 0.0, 0.0, 0.0, -STEP_GRIPPER, 20),
        ("raise", 0.0, 0.0, STEP_Z_M, 0.0, 12),
        ("slide left (+Y)", 0.0, STEP_XY_M, 0.0, 0.0, 16),
        ("slide right (-Y)", 0.0, -STEP_XY_M, 0.0, 0.0, 32),
        ("back to centre", 0.0, STEP_XY_M, 0.0, 0.0, 16),
        ("pull in (-X)", -STEP_XY_M, 0.0, 0.0, 0.0, 16),
        ("home", None, None, None, None, 1),
    ]

    dt = 1.0 / LOOP_HZ
    try:
        while True:
            for label, dx, dy, dz, dgrip, steps in script:
                print(f"  -> {label}")
                for _ in range(steps):
                    if dx is None:
                        ctrl.reset()
                    else:
                        ctrl.apply_command(dx, dy, dz, dgrip)
                    arm.update(ctrl.render_joint_values(), ctrl.target_pos)
                    _, ori = ctrl.residuals()
                    if ori > MAX_ORIENTATION_RESIDUAL_DEG:
                        print(f"     (warning: gripper tilted {ori:.1f} deg off vertical)")
                    time.sleep(dt)
                time.sleep(0.3)
    except KeyboardInterrupt:
        print("\nStopped.")


# ========================================================================================
# Browser-native control (recommended): a tiny stdlib HTTP server serves a SINGLE three.js
# page that renders the real arm meshes AND captures the keyboard. Because everything lives in
# one document, arrow keys hit our handler (preventDefault) instead of scrolling an embedded
# viewer, and no terminal focus / pynput / OS Input-Monitoring permission is needed.
#   /            -> the page          /meshes/<f>  -> STL bytes
#   /scene       -> static visuals    /frames      -> live per-link transforms (+ status)
#   /cmd?k=<key> -> queue a command
# ========================================================================================
CONTROL_HTML = """<!doctype html><html><head><meta charset="utf-8"><title>TOM teleop</title>
<style>
 html,body{margin:0;height:100%;overflow:hidden;background:#0e1014;color:#e6e6e6;font-family:system-ui,sans-serif}
 #app{display:flex;height:100vh}
 #side{width:228px;flex:none;padding:16px;box-sizing:border-box;background:#171a20}
 #view{flex:1;position:relative}
 canvas{display:block}
 h1{font-size:15px;margin:0 0 6px}
 .hint{font-size:12px;color:#9aa0a6;margin:8px 0;line-height:1.4}
 .grid{display:grid;grid-template-columns:repeat(3,1fr);gap:6px;margin:6px 0}
 button{padding:11px 0;font-size:13px;border:0;border-radius:6px;background:#2c313a;color:#e6e6e6;cursor:pointer}
 button:active{background:#37c871;color:#06210f}
 #status{font-size:12px;color:#37c871;margin-top:10px;white-space:pre-line;font-variant-numeric:tabular-nums}
</style>
<script type="importmap">
{ "imports": {
  "three": "https://unpkg.com/three@0.161.0/build/three.module.js",
  "three/addons/": "https://unpkg.com/three@0.161.0/examples/jsm/"
} }
</script>
<!-- livekit-client (classic global) for the OPTIONAL simulated camera feeds; if it fails to load the
     page still works as a local viewer. -->
<script src="https://cdn.jsdelivr.net/npm/livekit-client/dist/livekit-client.umd.min.js"></script>
</head><body>
<div id="app">
 <div id="side" style="overflow-y:auto">
  <h1>TOM &mdash; dual-arm teleop</h1>
  <div class="hint">Keyboard works anywhere on this page. Drag to orbit, scroll to zoom.
   Both grippers always face the ground; the chess set is for show.</div>
  <div style="font-weight:600;margin-top:8px">Arm 1 &mdash; arrows</div>
  <div class="grid">
   <div></div><button data-k="up">&#9650; away</button><div></div>
   <button data-k="left">&#9664; left</button><button data-k="down">&#9660; near</button><button data-k="right">right &#9654;</button>
  </div>
  <div class="grid">
   <button data-k="'">Z up &#39;</button><button data-k=".">Z down .</button><button data-k="h">home H</button>
   <button data-k="[">grip open [</button><button data-k="]">grip close ]</button><div></div>
  </div>
  <div class="grid">
   <button data-k="rest1">rest</button><button data-k="save1">save pose</button><button data-k="cust1">go custom</button>
  </div>
  <div id="arm2ui">
  <div style="font-weight:600;margin-top:10px">Arm 2 &mdash; WASD</div>
  <div class="grid">
   <div></div><button data-k="w">W away</button><div></div>
   <button data-k="a">A left</button><button data-k="s">S near</button><button data-k="d">D right</button>
  </div>
  <div class="grid">
   <button data-k="q">Z up Q</button><button data-k="e">Z down E</button><button data-k="b">home B</button>
   <button data-k="r">grip open R</button><button data-k="f">grip close F</button><div></div>
  </div>
  <div class="grid">
   <button data-k="rest2">rest</button><button data-k="save2">save pose</button><button data-k="cust2">go custom</button>
  </div>
  </div>
  <div id="status">loading 3D&hellip;</div>
  <div style="font-weight:600;margin-top:14px">Remote camera feeds (sim)</div>
  <div class="hint">Publish simulated wrist/top/side views over LiveKit so a remote operator can watch
   without real cameras. Optional &mdash; does not affect local control.</div>
  <button id="cambtn" style="width:100%;margin-top:4px">Stream cameras</button>
  <div id="camstatus" style="font-size:12px;color:#9aa0a6;margin-top:6px">off</div>
  <div id="campreviews" style="display:flex;gap:4px;flex-wrap:wrap;margin-top:6px"></div>
  <div style="font-weight:600;margin-top:14px">Joint control (verify limits)</div>
  <div class="hint">Set joints directly to their limits (bypasses IK). Controls
   <select id="jarm" style="background:#2c313a;color:#e6e6e6;border:0;border-radius:4px;padding:2px">
   <option value="arm1">Arm 1</option><option value="arm2">Arm 2</option></select></div>
  <div id="jsliders"></div>
 </div>
 <div id="view"></div>
</div>
<script type="module">
import * as THREE from 'three';
import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
import { STLLoader } from 'three/addons/loaders/STLLoader.js';
import { GLTFLoader } from 'three/addons/loaders/GLTFLoader.js';

const view=document.getElementById('view'), statusEl=document.getElementById('status');
THREE.Object3D.DEFAULT_UP.set(0,0,1);                       // URDF is Z-up
const renderer=new THREE.WebGLRenderer({antialias:true});
renderer.setPixelRatio(devicePixelRatio);
renderer.setSize(view.clientWidth, view.clientHeight);
view.appendChild(renderer.domElement);
const scene=new THREE.Scene(); scene.background=new THREE.Color(0xf1f3f5);  // light bg (overridden by /scene)
const camera=new THREE.PerspectiveCamera(45, view.clientWidth/view.clientHeight, 0.01, 50);
camera.up.set(0,0,1);
const controls=new OrbitControls(camera, renderer.domElement);
controls.enablePan=false;                                   // no key listeners -> arrows never move the camera
scene.add(new THREE.HemisphereLight(0xffffff,0x303038,1.15));
const dl=new THREE.DirectionalLight(0xffffff,1.6); dl.position.set(-0.5,-0.7,1.2); scene.add(dl);

// Ground reference: a 3D grid on the floor plane (just under the table). THREE.GridHelper is built
// in the XZ plane for a Y-up world, so rotate it into XY for our Z-up world.
(function addGrid(){
  const grid=new THREE.GridHelper(4, 40, 0x94a3b8, 0xcbd5e1);   // size 4 m, 40 divisions (10 cm cells)
  grid.rotation.x=Math.PI/2;                                    // XZ -> XY (lie flat in Z-up)
  grid.position.set(0.285, 0.0, -0.05);                         // centred on the table, at floor level
  scene.add(grid);
})();

const arms={};                                              // armId -> {groups:{link:group}}
const m4=a=>{const m=new THREE.Matrix4(); m.set(...a); return m;};   // a is row-major

function loadChess(ch){
  new GLTFLoader().load(ch.url, gltf=>{
    const o=gltf.scene;
    o.rotation.x=Math.PI/2;                                 // glTF Y-up -> our Z-up world
    if(ch.yaw) o.rotateOnWorldAxis(new THREE.Vector3(0,0,1), ch.yaw);  // face armies at the arms
    o.updateMatrixWorld(true);
    let box=new THREE.Box3().setFromObject(o), size=new THREE.Vector3(); box.getSize(size);
    o.scale.setScalar(ch.fit/Math.max(size.x,size.y));      // fit footprint to ch.fit metres
    o.updateMatrixWorld(true);
    box=new THREE.Box3().setFromObject(o); const c=new THREE.Vector3(); box.getCenter(c);
    o.position.set(ch.center[0]-c.x, ch.center[1]-c.y, -box.min.z);  // centre on table, sit at z=0
    scene.add(o);
  }, undefined, err=>console.warn('chess load failed:', err));
}

fetch('/scene').then(r=>r.json()).then(sc=>{
  if(sc.bg!=null) scene.background=new THREE.Color(sc.bg);
  const t=sc.table;
  // Circular table: a thin cylinder. CylinderGeometry's axis is +Y, so rotate it to stand on Z.
  const tableMesh=new THREE.Mesh(new THREE.CylinderGeometry(t.radius,t.radius,t.thickness,64),
    new THREE.MeshStandardMaterial({color:0x14161b,roughness:0.95,metalness:0.0}));
  tableMesh.rotation.x=Math.PI/2; tableMesh.position.set(...t.center); scene.add(tableMesh);
  if(sc.chess) loadChess(sc.chess);
  const loader=new STLLoader();
  for(const arm of sc.arms){
    const root=new THREE.Group(); root.matrixAutoUpdate=false;
    root.matrix.copy(m4(arm.world)); root.matrixWorldNeedsUpdate=true; scene.add(root);
    const g={};
    for(const link of sc.links){
      const lg=new THREE.Group(); lg.matrixAutoUpdate=false; root.add(lg); g[link.link]=lg;
      for(const v of link.visuals){
        loader.load('/meshes/'+v.mesh, geo=>{
          const mesh=new THREE.Mesh(geo, new THREE.MeshStandardMaterial(
            {color:v.color, roughness:0.55, metalness:0.25}));
          mesh.matrixAutoUpdate=false; mesh.matrix.copy(m4(v.matrix)); lg.add(mesh);
        });
      }
    }
    arms[arm.id]={groups:g, root};
  }
  camera.position.set(...sc.camera.eye); controls.target.set(...sc.camera.target); controls.update();
  statusEl.textContent='ready';
  buildJointSliders(sc.joint_order, sc.joint_limits);
  if(sc.single_arm){                                   // hide all arm-2 UI in single-arm mode
    const a2=document.getElementById('arm2ui'); if(a2) a2.style.display='none';
    const opt=document.querySelector('#jarm option[value="arm2"]'); if(opt) opt.remove();
  }
  if(sc.include_gripper===false){                      // no gripper attached -> hide its controls
    document.querySelectorAll('button').forEach(b=>{ if(["[","]","r","f"].includes(b.dataset.k)) b.style.display='none'; });
  }
  try{ setupCamViews();                                // build the optional remote camera views
    if(new URLSearchParams(location.search).get('stream')==='1') toggleStream();  // auto-publish
  }catch(e){ console.warn('cam views setup failed', e); }
});

// ---- joint sliders (verify limits): set each joint directly via /joint ----
let jointOrder=[];
const jarm=()=>document.getElementById('jarm').value;
function buildJointSliders(order, limits){
  jointOrder=order||[];
  const box=document.getElementById('jsliders'); box.innerHTML='';
  for(const n of jointOrder){
    const [lo,hi]=limits[n];
    const row=document.createElement('div'); row.style.margin='5px 0';
    const lab=document.createElement('div'); lab.id='jl_'+n; lab.style.fontSize='11px';
    lab.style.color='#9aa0a6'; lab.textContent=`${n}: 0  [${lo}…${hi}]`;
    const inp=document.createElement('input'); inp.type='range'; inp.min=lo; inp.max=hi;
    inp.step=1; inp.value=0; inp.style.width='100%'; inp.dataset.j=n;
    inp.addEventListener('input',()=>{ lab.textContent=`${n}: ${(+inp.value).toFixed(0)}  [${lo}…${hi}]`;
      fetch(`/joint?arm=${jarm()}&j=${n}&v=${inp.value}`).catch(()=>{}); });
    row.appendChild(lab); row.appendChild(inp); box.appendChild(row);
  }
}

const fmt=s=>`x=${s.x.toFixed(3)} y=${s.y.toFixed(3)} z=${s.z.toFixed(3)}  grip ${s.grip.toFixed(0)}%`;
async function pollFrames(){
  try{ const f=await (await fetch('/frames')).json();
    for(const id in arms){ const fr=f[id]; if(!fr) continue; const g=arms[id].groups;
      for(const k in fr){ if(g[k]){ g[k].matrix.copy(m4(fr[k])); g[k].matrixWorldNeedsUpdate=true; } } }
    const s=f.__status;
    if(s){
      statusEl.textContent='Arm 1  '+fmt(s.arm1)+(s.arm2?('\\nArm 2  '+fmt(s.arm2)):'');
      const js=s[jarm()] && s[jarm()].joints;            // live joints of the selected arm
      if(js) for(const n of jointOrder){
        const inp=document.querySelector(`#jsliders input[data-j="${n}"]`); if(!inp) continue;
        const lab=document.getElementById('jl_'+n);
        const v=js[n]; if(lab) lab.textContent=`${n}: ${v.toFixed(0)}  [${inp.min}…${inp.max}]`;
        if(document.activeElement!==inp) inp.value=v;     // don't fight a slider being dragged
      }
    }
  }catch(e){}
}
setInterval(pollFrames, 50);

const send=k=>fetch('/cmd?k='+encodeURIComponent(k)).catch(()=>{});
const map={ArrowUp:'up',ArrowDown:'down',ArrowLeft:'left',ArrowRight:'right',Escape:'esc'};
const letters=new Set(['h','w','a','s','d','q','e','r','f','b',"'",'.','[',']']);
addEventListener('keydown',e=>{
  let k=map[e.key]; if(!k){const lk=(e.key||'').toLowerCase(); if(letters.has(lk))k=lk;}
  if(k){ e.preventDefault(); send(k); }
}, {passive:false});
document.querySelectorAll('button').forEach(b=>b.addEventListener('click',()=>send(b.dataset.k)));
addEventListener('resize',()=>{ renderer.setSize(view.clientWidth,view.clientHeight);
  camera.aspect=view.clientWidth/view.clientHeight; camera.updateProjectionMatrix(); });

// ---- OPTIONAL simulated camera feeds: extra views of the SAME scene, published over LiveKit so a
//      remote operator can watch without real cameras. Each view is its own small canvas + renderer;
//      publishTrack(canvas.captureStream(...)) sends it as a named video track. Fully optional. ----
const CAM_W=480, CAM_H=360; const camViews=[]; let lkRoom=null;
const camStatus=t=>{ const e=document.getElementById('camstatus'); if(e) e.textContent=t; };
function makeCamView(name, fov){
  const cv=document.createElement('canvas'); cv.width=CAM_W; cv.height=CAM_H; cv.title=name;
  cv.style.width='110px'; cv.style.height='82px'; cv.style.background='#000';
  document.getElementById('campreviews').appendChild(cv);          // tiny live preview in the panel
  const r=new THREE.WebGLRenderer({canvas:cv, antialias:true}); r.setSize(CAM_W,CAM_H,false);
  const cam=new THREE.PerspectiveCamera(fov||52, CAM_W/CAM_H, 0.01, 50); cam.up.set(0,0,1);
  return {name, canvas:cv, renderer:r, cam, place:null};
}
function gripperLinkOf(arm){
  const g=arm.groups||{};
  for(const k of Object.keys(g)) if(/gripper.*base|gripper|tool|tip/i.test(k)) return g[k];
  const ks=Object.keys(g); return ks.length?g[ks[ks.length-1]]:null;   // fallback: last link
}
function setupCamViews(){
  if(camViews.length || !Object.keys(arms).length) return;
  const top=makeCamView('top');                                        // bird's-eye over the table
  top.place=()=>{ top.cam.position.set(0.285,0.0,1.35); top.cam.up.set(0,1,0); top.cam.lookAt(0.285,0.0,0.0); };
  const side=makeCamView('side');                                      // along -Y
  side.place=()=>{ side.cam.position.set(0.285,-1.15,0.5); side.cam.up.set(0,0,1); side.cam.lookAt(0.285,0.0,0.30); };
  camViews.push(top, side);
  // TOOL_APPROACH_AXIS (gripper_base frame): the unit direction the fingers POINT — mirrors the Python
  // control core so the wrist cam looks exactly where the gripper reaches (~ local -Z, tiny +X tilt).
  const APPROACH_X=0.02070, APPROACH_Z=-0.99979;
  for(const id of Object.keys(arms)){                                  // egocentric wrist cam per arm
    const name = id==='arm1' ? 'wrist_1' : 'wrist_2';
    const w=makeCamView(name, 72);               // wide FOV, like a real wrist cam
    const base=gripperLinkOf(arms[id]);          // gripper_base_link = the camera-stand mount frame
    // Mount the camera AT THE END of the end-effector (the claw tips, ~TIP_M along the reach from the
    // gripper_base) looking OUTWARD along the reach — where the gripper points. This depends only on
    // the reach axis (TOOL_APPROACH_AXIS), so there's no left/right ambiguity and the view never falls
    // back toward the arm.
    const TIP_M=0.04;                            // end-of-effector offset along reach (claw tips ~4 cm)
    w.place=()=>{
      if(!base) return;
      base.updateWorldMatrix(true,false); const m=base.matrixWorld;
      const basePos=new THREE.Vector3().setFromMatrixPosition(m);
      const colX=new THREE.Vector3().setFromMatrixColumn(m,0).normalize();        // local X (+X toward arm)
      const colZ=new THREE.Vector3().setFromMatrixColumn(m,2).normalize();        // local Z = up the wrist
      const reach=new THREE.Vector3()                                             // fingers-point dir, world
        .addScaledVector(colX, APPROACH_X).addScaledVector(colZ, APPROACH_Z).normalize();
      const tip=basePos.clone().addScaledVector(reach, TIP_M);                    // end of the end-effector
      w.cam.up.copy(colX).multiplyScalar(-1);                                     // outward side = image up
      w.cam.position.copy(tip);
      w.cam.lookAt(tip.clone().addScaledVector(reach, 0.30));                     // look OUTWARD along reach
    };
    camViews.push(w);
  }
}
function renderCamViews(){ for(const v of camViews){ try{ if(v.place)v.place(); v.renderer.render(scene,v.cam); }catch(e){} } }
async function toggleStream(){
  if(lkRoom){ try{ await lkRoom.disconnect(); }catch(e){} lkRoom=null; camStatus('off');
    document.getElementById('cambtn').textContent='Stream cameras'; return; }
  if(typeof LivekitClient==='undefined'){ camStatus('livekit-client not loaded (needs internet)'); return; }
  setupCamViews();
  if(!camViews.length){ camStatus('scene not ready yet'); return; }
  let cfg; try{ cfg=await (await fetch('/token?identity=sim-cams')).json(); }catch(e){ camStatus('token fetch failed'); return; }
  if(cfg.error){ camStatus('token: '+cfg.error+' (run with --livekit-* set)'); return; }
  const room=new LivekitClient.Room();
  try{ await room.connect(cfg.url, cfg.token); }catch(e){ camStatus('connect failed: '+(e&&e.message)); return; }
  lkRoom=room;
  for(const v of camViews){
    try{ const t=v.canvas.captureStream(15).getVideoTracks()[0];
      const lt=new LivekitClient.LocalVideoTrack(t);
      await room.localParticipant.publishTrack(lt, {name:v.name, source:LivekitClient.Track.Source.Camera});
    }catch(e){ console.warn('publish '+v.name+' failed', e); }
  }
  camStatus('streaming '+cfg.room+': '+camViews.map(v=>v.name).join(', '));
  document.getElementById('cambtn').textContent='Stop streaming';
}
document.getElementById('cambtn').addEventListener('click', toggleStream);

(function loop(){
  controls.update(); renderer.render(scene,camera); renderCamViews(); requestAnimationFrame(loop);
})();
</script></body></html>"""


class RemoteTargetServer:
    """UDP server for remote teleop. Bind to 127.0.0.1 in production and feed it only from a local
    bridge (e.g. ``-m tomctl.livekit.bridge``); the bridge is the authenticated edge. This server
    itself does NO authentication, so never bind it to a routable interface for remote use.

    Accepts two message shapes (both keep the deadman alive); latest-wins, so it tolerates jitter,
    reordering and packet loss (a dropped packet = a slightly stale state the next one corrects):

    1. ABSOLUTE tip targets (legacy ``tools/latency_probe.py`` / sim web):
         {"seq": 12, "t": 1234.5, "targets": {"arm1": {"x":..,"y":..,"z":..,"gripper":..}}}
       Fed straight to ``command_absolute``. Non-finite values are rejected; a strictly-older seq
       that arrives within the same live session is dropped (reorder/replay guard).

    2. JOG INTENT (the keyboard-over-LiveKit path): the SET of currently-held movement keys, translated
       on THIS side into the SAME ``command_delta`` the local keyboard uses (remote == local):
         {"seq": 30, "t": .., "type": "intent", "arm": "arm1", "deadman": true,
          "held": ["up","left","."],            # ARM1_BIND jog tokens, applied while held
          "edges": [{"id": 88, "action": "session_start"}]}   # one-shot, de-duped by id
       ``deadman`` must be true for held-key jog to apply (positive interlock). The only edges are the
       remote-only session lifecycle ``session_start`` (torque on + home) / ``session_end`` (rest +
       torque off, deferred until settled), injected by the bridge — the operator sends no edges.
       ``arm`` selects the target controller (default "arm1" if omitted); each arm keeps fully
       independent held/deadman/edge state, so two operators can drive arm1 + arm2 at once.

    A bare {"hb": true} datagram is a heartbeat (keeps the deadman alive with no motion change).

    DEADMAN: if no datagram arrives within ``DEADMAN_S``, motion is frozen (hold pose)."""

    DEADMAN_S = 0.4  # hold position if no packet/heartbeat arrives within this long
    _SESSION_GAP_S = 1.0  # a gap longer than this resets the per-arm seq baseline (new session)

    # Held jog tokens -> (intent field, signed speed). Speeds match the Joy-Con defaults so remote
    # held-key jog feels like the gamepad: 0.084 m/s XY, 0.056 m/s Z, 56 gripper-units/s — a 30%
    # slowdown from the original 0.12/0.08/80. Only the target VELOCITY is lower; the accel ramp
    # (CART_MAX_ACCEL_MPS2) is unchanged, so smoothness is unaffected. Local keyboard jog derives its
    # magnitudes from this map (see real.py _JOG_MAG), so both local and remote slow together.
    JOG_VELOCITY = {
        "up": ("dx", +0.084), "down": ("dx", -0.084),
        "left": ("dy", +0.084), "right": ("dy", -0.084),
        "'": ("dz", +0.056), ".": ("dz", -0.056),
        "[": ("dgrip", +56.0), "]": ("dgrip", -56.0),
    }
    # The ONLY edges are the REMOTE queue's automatic lifecycle (no local-keyboard equivalent): start =
    # torque ON + go HOME; end = go REST + torque OFF (deferred until the rest move settles). Injected by
    # the LiveKit bridge at session boundaries; the operator and the local teleop never send edges.
    EDGE_ACTIONS = frozenset({"session_start", "session_end"})

    def __init__(self, port: int, host: str = "127.0.0.1",
                 torque_cb: "Callable[[str, bool], None] | None" = None) -> None:
        import socket
        import threading

        # Loopback by default: this server is UNAUTHENTICATED, so it must only ever be fed by a local
        # bridge (the authenticated edge). Pass host="0.0.0.0" explicitly ONLY for a trusted LAN test
        # with tools/latency_probe.py — never for a network reachable by untrusted hosts.
        self.port = port
        # Optional per-arm torque control, e.g. robot.bus.enable/disable_torque on the real arm. None on
        # the sim (no motors), so session_start/session_end only move the pose. Signature: (arm_id, on).
        self._torque_cb = torque_cb
        self._release_pending: set[str] = set()  # arms parking to rest, awaiting torque-off when settled
        self._pending_start: set[str] = set()  # session_start deferred until other arms finish parking (1-at-a-time)
        self._lock = threading.Lock()
        self._latest: dict[str, tuple[int, tuple[float, float, float, float]]] = {}
        self._seq_t: dict[str, float] = {}  # monotonic time of the last accepted seq per arm
        self._last_rx: float | None = None
        # PER-ARM jog-intent state, keyed by the datagram's ``arm`` field (default "arm1"). Two
        # operators driving arm1 and arm2 concurrently keep fully independent held/deadman/estop/edge
        # state. deadman defaults False (fail-safe positive interlock): a frame must explicitly assert
        # deadman=true for any jog to apply. See ``_arm_intent`` for the per-arm record shape.
        self._intent: dict[str, dict] = {}
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._sock.bind((host, port))
        self._sock.settimeout(0.2)
        self._stop = False
        self._thread = threading.Thread(target=self._recv_loop, daemon=True)
        self._thread.start()

    def _arm_intent(self, aid: str) -> dict:
        """Lazily create the per-arm jog-intent record (call under self._lock)."""
        st = self._intent.get(aid)
        if st is None:
            st = {"held": set(), "deadman": False, "intent_seq": -1, "intent_t": 0.0,
                  "pending_edges": [], "seen_edge_ids": set()}
            self._intent[aid] = st
        return st

    def _recv_loop(self) -> None:
        import json
        import math
        import socket
        import time

        while not self._stop:
            try:
                data, _addr = self._sock.recvfrom(8192)
            except socket.timeout:
                continue
            except OSError:
                break
            try:
                msg = json.loads(data.decode())
            except Exception:
                continue
            if not isinstance(msg, dict):
                continue
            now = time.monotonic()
            raw_seq = msg.get("seq", 0)
            seq = int(raw_seq) if isinstance(raw_seq, (int, float)) else 0
            with self._lock:
                self._last_rx = now
                # --- absolute targets (legacy) ---
                targets = msg.get("targets")
                if isinstance(targets, dict):
                    for aid, t in targets.items():
                        if not isinstance(t, dict):
                            continue
                        try:
                            tup = (float(t["x"]), float(t["y"]), float(t["z"]), float(t.get("gripper", 0.0)))
                        except (KeyError, TypeError, ValueError):
                            continue
                        if not all(math.isfinite(v) for v in tup):
                            continue  # reject NaN/inf before it can poison the IK/slew
                        prev = self._latest.get(aid)
                        # Drop a strictly-older seq that arrives inside the same live session
                        # (reorder/replay). A gap > _SESSION_GAP_S is a new session -> accept + reset.
                        if prev is not None and seq < prev[0] and (now - self._seq_t.get(aid, 0.0)) < self._SESSION_GAP_S:
                            continue
                        self._latest[aid] = (seq, tup)
                        self._seq_t[aid] = now
                # --- jog intent (new), routed to the datagram's arm (default arm1) ---
                if "held" in msg or "edges" in msg or msg.get("type") == "intent":
                    aid = msg.get("arm", "arm1")
                    if not isinstance(aid, str):
                        aid = "arm1"
                    st = self._arm_intent(aid)
                    # A gap (e.g. the bridge restarted, so its seq counter restarts low) is a new
                    # session: reset the seq baseline and DROP the held set so stale motion can't
                    # resume. Otherwise a low new seq would be ignored and held would stay frozen.
                    if now - st["intent_t"] > self._SESSION_GAP_S:
                        st["intent_seq"] = -1
                        st["held"] = set()
                        st["deadman"] = False
                    st["intent_t"] = now
                    if seq >= st["intent_seq"]:  # latest-wins held set
                        st["intent_seq"] = seq
                        held = msg.get("held", [])
                        if isinstance(held, list):
                            st["held"] = {k for k in held if k in self.JOG_VELOCITY}
                        st["deadman"] = bool(msg.get("deadman", False))
                    edges = msg.get("edges", [])
                    if isinstance(edges, list):
                        for e in edges:
                            if not isinstance(e, dict):
                                continue
                            try:
                                eid, action = int(e["id"]), str(e["action"])
                            except (KeyError, TypeError, ValueError):
                                continue
                            if action in self.EDGE_ACTIONS and eid not in st["seen_edge_ids"]:
                                st["seen_edge_ids"].add(eid)
                                st["pending_edges"].append((eid, action))
                    if len(st["seen_edge_ids"]) > 4096:  # bound memory; ids are monotonic per session
                        st["seen_edge_ids"] = set(sorted(st["seen_edge_ids"])[-1024:])

    def apply(self, arms: dict, dt: float) -> None:
        """Push the latest received intent into the controllers (called each control tick)."""
        import time

        now = time.monotonic()
        with self._lock:
            last_rx = self._last_rx
            latest = dict(self._latest)
            # snapshot + drain each arm's pending edges atomically
            snaps = {}
            for aid, st in self._intent.items():
                snaps[aid] = {"held": set(st["held"]), "deadman": st["deadman"],
                              "intent_t": st["intent_t"], "edges": st["pending_edges"]}
                st["pending_edges"] = []
        if last_rx is None:
            return  # no remote traffic yet -> leave the keyboard in charge

        def _hold_one(c) -> None:
            c.commanded_pos = c.target_pos.copy()

        # --- legacy ABSOLUTE targets (tools/latency_probe.py / sim web): keyed by arm, global deadman.
        # Mutually exclusive in practice with jog-intent (a different sender type), so it does not
        # interact with the per-arm estop latch below.
        abs_stale = now - last_rx > self.DEADMAN_S
        for aid, (_seq, tgt) in latest.items():
            c = arms.get(aid)
            if c is None:
                continue
            if abs_stale:
                _hold_one(c)
            else:
                c.command_absolute(*tgt)

        # --- PER-ARM jog intent (keyboard-over-LiveKit). Each arm is independent. ---
        for aid, snap in snaps.items():
            c = arms.get(aid)
            if c is None:
                continue
            # 1) one-shot lifecycle edges (already de-duped) — the bridge injects these automatically
            #    at the start/end of a remote session; the operator sends no edges of their own.
            for _eid, action in snap["edges"]:
                if action == "session_start":
                    # session begins: power the arm and raise to HOME — but DEFER it while ANOTHER arm is
                    # still parking to rest (torque not yet released), so at most ONE arm is ever energized
                    # (home/rest transitions never overlap; no double current draw). The gated block after
                    # this loop fires it. A re-take of the SAME arm mid-park cancels its own park (discard
                    # below) and, with no other arm resting, homes on this same tick.
                    self._release_pending.discard(aid)  # re-taken: cancel its own pending park-release
                    self._pending_start.add(aid)
                elif action == "session_end":
                    # session ends: park to REST, then release torque once the move settles.
                    self._pending_start.discard(aid)  # a deferred start that never fired -> drop it
                    c.go_rest()
                    if self._torque_cb is not None:
                        self._release_pending.add(aid)
            # torque-off is DEFERRED until the rest move finishes, so the arm settles supported.
            if aid in self._release_pending and not c.is_settling():
                if self._torque_cb is not None:
                    self._torque_cb(aid, False)
                self._release_pending.discard(aid)
            # 2) deadman TIMEOUT: this arm's link went silent -> freeze (safety backstop).
            if now - snap["intent_t"] > self.DEADMAN_S:
                _hold_one(c)
                continue
            # 3) held-key jog — needs the positive deadman interlock (the frame must assert deadman).
            #    SUPPRESSED until the arm has finished its session-start HOME: command_delta cancels the
            #    running path and takes over from wherever the arm currently is, so an eager operator
            #    holding a key as their turn begins would cancel the raise-to-home and drive from the REST
            #    pose. Gate on BOTH is_settling() (the home/rest move is running) AND a pending start (this
            #    tick's session_start, or a start still DEFERRED behind another arm's park — not homed yet).
            if snap["deadman"] and snap["held"] and not c.is_settling() and aid not in self._pending_start:
                d = {"dx": 0.0, "dy": 0.0, "dz": 0.0, "dgrip": 0.0}
                for k in snap["held"]:
                    field, speed = self.JOG_VELOCITY[k]
                    d[field] += speed * dt
                if any(d.values()):
                    c.command_delta(d["dx"], d["dy"], d["dz"], d["dgrip"])

        # Fire any DEFERRED session_start now that no OTHER arm is still parking to rest — enforcing
        # one-arm-at-a-time: an arm only powers up + homes once the previous arm has fully settled at
        # rest and released torque. Runs after the whole per-arm loop so session_end (which marks the
        # other arm as parking) has already been applied this tick, regardless of arm iteration order.
        for said in list(self._pending_start):
            if any(a != said for a in self._release_pending):
                continue  # another arm still parking -> keep this start waiting
            sc = arms.get(said)
            if sc is not None:
                if self._torque_cb is not None:
                    self._torque_cb(said, True)
                sc.go_home()
            self._pending_start.discard(said)

    def close(self) -> None:
        self._stop = True
        try:
            self._sock.close()
        except Exception:
            pass


def run_web(
    model: UrdfModel,
    port: int = 7010,
    remote_port: int | None = None,
    single_arm: bool = False,
    include_gripper: bool | None = None,
    gamepad: bool = False,
    gamepad_profile: str | None = None,
    livekit_url: str = "ws://localhost:7880",
    livekit_room: str = "tom-cams",
    livekit_key: str = "devkey",
    livekit_secret: str = "secret",
    web_host: str = "127.0.0.1",
) -> None:
    import http.server
    import json
    import queue
    import socketserver
    import threading
    import urllib.parse
    import webbrowser

    if include_gripper is not None:  # explicit flag overrides any calibration-loaded value
        globals()["INCLUDE_GRIPPER"] = include_gripper

    remote = RemoteTargetServer(remote_port) if remote_port else None

    # Two arms: arm1 at the world origin (reaches +X), arm2 at the +X edge rotated 180 deg about
    # Z (reaches -X). Each controller solves in its OWN frame; the world placement is applied by
    # the renderer's parent group, so the same FK works for both.
    arms = {"arm1": CartesianController(model)}
    if not single_arm:
        arms["arm2"] = CartesianController(model)

    # Optional Joy-Con teleop: one controller per arm (arm1, and arm2 unless single-arm). Built BEFORE
    # the browser/server starts so a missing mapping errors cleanly. Polled on THIS (main) thread each
    # tick — macOS only delivers SDL joystick events when pumped on main — and routed into the SAME
    # command_delta/named-pose intents as the keyboard via tomctl.gamepad.
    pads = []
    pygame_mod = None
    gi = None
    if gamepad:
        from tomctl import gamepad as gi  # noqa: PLC0415  the shared gamepad mapping module

        roles = ["arm1"] if single_arm else list(arms)
        print("Joy-Con teleop: pairing controllers to arms...")
        pads = gi.build_gamepad_set(roles, profile_path=gamepad_profile or gi.DEFAULT_PROFILE_PATH)
        pygame_mod = gi._require_pygame()

    def world_placement(yaw_deg: float, x: float) -> NDArray[np.float64]:
        a = np.deg2rad(yaw_deg)
        w = np.eye(4)
        w[:3, :3] = np.array([[np.cos(a), -np.sin(a), 0], [np.sin(a), np.cos(a), 0], [0, 0, 1]])
        w[0, 3] = x
        return w

    # Fixed circular table: the two arms clamp flush to opposite edges of the 90 cm disk along the
    # X (arm-separation) axis, board centred between them. CLAMP_INSET_M (module constant) is the base
    # origin's inboard offset from the clamped edge (ARM1_BASE_X = (table_center - diameter/2) + inset).
    arm_world: dict[str, NDArray[np.float64]] = {}
    cx = TABLE_CENTER[0]
    half = TABLE_DIAMETER / 2.0
    arm_world["arm1"] = world_placement(0.0, (cx - half) + CLAMP_INSET_M)
    arm_world["arm2"] = world_placement(180.0, (cx + half) - CLAMP_INSET_M)
    assets_dir = model.urdf_path.parent.resolve()  # assets/ (chess bundle lives here)
    meshes_dir = assets_dir / "meshes"
    cmd_q: queue.Queue[str] = queue.Queue()
    joint_q: queue.Queue[tuple[str, str, float]] = queue.Queue()  # (arm_id, joint_name, value_deg)

    # Joint sliders (verify limits): name -> [lo, hi], in base->tip order plus the gripper.
    joint_order = [*IK_ACTIVE_JOINTS] + (["gripper"] if INCLUDE_GRIPPER else [])
    slider_limits = {n: list(JOINT_LIMITS_DEG[n]) for n in IK_ACTIVE_JOINTS}
    if INCLUDE_GRIPPER:
        slider_limits["gripper"] = list(GRIPPER_LIMITS)

    # Static scene: shared link visuals (mesh file, colour, baked visual-origin @ scale matrix),
    # the two arms' world placements, the table, the chess set, and the camera. Matrices are
    # row-major (THREE.Matrix4.set order).
    def scene_payload() -> dict:
        links = [
            {
                "link": link,
                "visuals": [
                    {
                        "mesh": v.mesh_path.name,
                        "color": v.color_hex,
                        "matrix": (v.origin @ np.diag([*v.scale, 1.0])).flatten().tolist(),
                    }
                    for v in visuals
                ],
            }
            for link, visuals in model.links.items()
            if visuals and (INCLUDE_GRIPPER or "gripper" not in link)  # hide gripper meshes if absent
        ]
        return {
            "links": links,
            "include_gripper": INCLUDE_GRIPPER,
            "arms": [{"id": aid, "world": arm_world[aid].flatten().tolist()} for aid in arms],
            "table": {
                "center": list(TABLE_CENTER),
                "shape": "circle",
                "radius": TABLE_DIAMETER / 2.0,
                "thickness": TABLE_THICKNESS,
            },
            "chess": {
                "url": CHESS_GLTF_URL,
                "fit": CHESS_FIT_SIZE,
                "center": list(TABLE_CENTER[:2]),
                "yaw": float(np.deg2rad(CHESS_YAW_DEG)),
            },
            "camera": {"eye": list(WEB_CAMERA_EYE), "target": list(WEB_CAMERA_TARGET)},
            "bg": WEB_BG_COLOR,
            "joint_order": joint_order,
            "joint_limits": slider_limits,
            "single_arm": single_arm,
        }

    def frames_payload() -> dict:
        out: dict = {}
        status: dict = {}
        for aid, ctrl in arms.items():
            frames = model.forward_kinematics_links(ctrl.render_joint_values())
            out[aid] = {link: t.flatten().tolist() for link, t in frames.items()}
            p = ctrl.target_pos
            status[aid] = {
                "x": float(p[0]),
                "y": float(p[1]),
                "z": float(p[2]),
                "grip": float(ctrl.gripper),
                "joints": ctrl.joint_state(),
            }
        out["__status"] = status
        return out

    def _send(handler, code, body=b"", ctype="text/plain"):
        handler.send_response(code)
        if body:
            handler.send_header("Content-Type", ctype)
            handler.send_header("Content-Length", str(len(body)))
        handler.end_headers()
        if body:
            handler.wfile.write(body)

    class Handler(http.server.BaseHTTPRequestHandler):
        def log_message(self, *a):  # silence per-request logging
            pass

        def do_GET(self):
            u = urllib.parse.urlparse(self.path)
            if u.path in ("/", "/index.html"):
                _send(self, 200, CONTROL_HTML.encode(), "text/html; charset=utf-8")
            elif u.path == "/scene":
                _send(self, 200, json.dumps(scene_payload()).encode(), "application/json")
            elif u.path == "/token":
                # Mint a PUBLISHER token so the browser page can stream its simulated camera views to
                # LiveKit. livekit is imported lazily, so the LOCAL sim never needs it; the route just
                # 503s if it's absent. Dev/LAN key/secret only.
                ident = urllib.parse.parse_qs(u.query).get("identity", ["sim-cams"])[0]
                try:
                    from livekit import api as _lkapi  # noqa: PLC0415  optional, remote-only
                    tok = (_lkapi.AccessToken(livekit_key, livekit_secret)
                           .with_identity(ident).with_name("TOM sim cameras")
                           .with_grants(_lkapi.VideoGrants(
                               room_join=True, room=livekit_room,
                               can_publish=True, can_subscribe=False, can_publish_data=False))
                           .to_jwt())
                    _send(self, 200, json.dumps(
                        {"url": livekit_url, "room": livekit_room, "token": tok}).encode(),
                        "application/json")
                except Exception as exc:  # noqa: BLE001
                    _send(self, 503, json.dumps({"error": f"livekit unavailable: {exc}"}).encode(),
                          "application/json")
            elif u.path == "/frames":
                _send(self, 200, json.dumps(frames_payload()).encode(), "application/json")
            elif u.path.startswith("/meshes/"):
                stl = meshes_dir / Path(u.path[len("/meshes/") :]).name  # basename only (no traversal)
                if stl.is_file():
                    _send(self, 200, stl.read_bytes(), "model/stl")
                else:
                    _send(self, 404)
            elif u.path.startswith("/assets/"):
                # Serve the vendored chess bundle (gltf + bin + textures). Keep relative subdirs
                # (textures/) but stay strictly inside assets_dir.
                rel = urllib.parse.unquote(u.path[len("/assets/") :])
                target = (assets_dir / rel).resolve()
                ctype = _CONTENT_TYPES.get(target.suffix.lower(), "application/octet-stream")
                if assets_dir in target.parents and target.is_file():
                    _send(self, 200, target.read_bytes(), ctype)
                else:
                    _send(self, 404)
            elif u.path == "/cmd":
                k = urllib.parse.parse_qs(u.query).get("k", [""])[0]
                if k:
                    cmd_q.put(k)
                _send(self, 204)
            elif u.path == "/joint":
                qd = urllib.parse.parse_qs(u.query)
                aid = qd.get("arm", ["arm1"])[0]
                jn = qd.get("j", [""])[0]
                try:
                    jv = float(qd.get("v", [""])[0])
                except ValueError:
                    jv = None
                if aid in arms and jn and jv is not None:
                    joint_q.put((aid, jn, jv))
                _send(self, 204)
            else:
                _send(self, 404)

    socketserver.ThreadingTCPServer.allow_reuse_address = True
    httpd = socketserver.ThreadingTCPServer((web_host, port), Handler)
    httpd.daemon_threads = True
    threading.Thread(target=httpd.serve_forever, daemon=True).start()

    url = f"http://127.0.0.1:{port}/"
    webbrowser.open(url)
    print("\nControl panel (drive it here with keyboard or buttons):", url)
    print("Arm 1: arrows -> XY | '/. -> up/down | [/] -> grip open/close | h -> home")
    if not single_arm:
        print("Arm 2: WASD  -> XY | q/e -> up/down | r/f -> grip open/close | b -> home")
    print("Poses (buttons): home / rest / save pose / go custom  (rest & custom always route via home)")
    if single_arm:
        print("Single-arm mode (arm 1 only).")
    print()
    if remote is not None:
        print(f"Remote teleop: accepting UDP targets on port {remote.port} (deadman {remote.DEADMAN_S}s)\n")
    if pads:
        print(
            f"Joy-Con teleop ({len(pads)} controller(s)): "
            "stick = XY | Z/grip buttons | home/rest/save/custom buttons | hold deadman to freeze\n"
        )

    period = 1.0 / LOOP_HZ
    last = time.monotonic()
    quit_now = False
    try:
        while not quit_now:
            loop_start = time.monotonic()
            dt = min(max(loop_start - last, 1e-3), 0.1)  # REAL elapsed time -> correct slew speed
            last = loop_start
            # Keyboard/button events update each arm's INTENT (commanded target) — they do not move
            # the arm directly; the control_step() below slews the working target toward it smoothly.
            while True:
                try:
                    key = cmd_q.get_nowait().lower()
                except queue.Empty:
                    break
                if key in ("q", "esc"):
                    quit_now = True
                elif not _apply_key_intent(key, arms["arm1"], ARM1_BIND) and "arm2" in arms:
                    _apply_key_intent(key, arms["arm2"], ARM2_BIND)
            while True:  # direct joint-slider sets (bypass IK)
                try:
                    aid, jn, jv = joint_q.get_nowait()
                except queue.Empty:
                    break
                arms[aid].set_joint(jn, jv)
            # Network (UDP) targets, if the remote server is enabled: latest-wins absolute targets,
            # with a deadman that freezes the intent if the link goes silent.
            if remote is not None:
                remote.apply(arms, dt)
            if pads:  # Joy-Con intent (per arm) — same command_delta/named-pose path as the keyboard
                pygame_mod.event.pump()
                for pad in pads:
                    cmd = pad.poll(dt)
                    if cmd.quit:
                        quit_now = True
                    if pad.role in arms:
                        gi.apply_command(cmd, arms[pad.role])
            for ctrl in arms.values():  # fixed-rate smooth execution
                ctrl.control_step(dt)
            time.sleep(max(0.0, period - (time.monotonic() - loop_start)))  # pace to LOOP_HZ
    except KeyboardInterrupt:
        pass
    finally:
        if remote is not None:
            remote.close()
        httpd.shutdown()
        print("\nStopped.")


CALIBRATION_KEYS = (  # JSON keys -> notes (documents the IK-control calibration schema)
    "joint_limits_deg",  # {joint: [lo, hi]} — the real arm's calibrated travel; the IK clamps to these
    "gripper_limits",  # [min, max] — the gripper's open<->closed travel (0..100); clamps the gripper command
    "home_pose_deg",  # {joint: deg} — HOME (ready) pose
    "rest_pose_deg",  # {joint: deg} — stowed REST pose (reached via home)
    "tool_length_m",  # gripper-tip offset along the approach axis (the controlled TCP)
    "approach_cap_deg",  # max gripper tilt off vertical a move may use (reach vs grasp)
    "floor_z_m",  # work-surface height (the target is never driven below this)
    "jog_frame_yaw_deg",  # operator "forward" (UP) heading in world deg; omit/null = auto from home
    "jog_flip_lr",  # true = mirror LEFT/RIGHT (left-handed world XY); else right-handed
    "gripper_invert",  # true = send (100 - command) to the gripper servo (servo-low = jaw open)
    "workspace_radius_m",  # table-disk radius (m) confining the tip during MANUAL control; null = off
    "workspace_center_xy",  # [x, y] table-disk centre in the arm's OWN IK frame (m); real arm ~[0, 0.381]

    "motion",  # {cart_max_speed_mps, cart_max_accel_mps2, joint_max_speed_dps,
    #            joint_path_speed_dps, joint_path_accel_dps2, gripper_speed_pps}
    "gripper",  # {servo_open_deg, servo_closed_deg} — command->servo endpoints (inverted mapping)
)


def load_calibration(source: str | Path | dict, expected_convention: str | None = None) -> dict:
    """Apply an IK-control calibration (a JSON path or an already-parsed dict) over the in-file
    defaults. This is the IK-control layer (home/rest/limits/gripper/motion/TCP) — it complements,
    and does NOT replace, lerobot's hardware (servo zero/range) calibration. Returns the config.

    Calibrations are CONVENTION-SPECIFIC: the sim's CAD URDF and the real arm's vendored URDF use
    different joint zero/axis conventions, so a home/rest/limit set captured on one is meaningless on
    the other. If ``expected_convention`` is given and the calibration declares a different
    ``convention``, it is NOT applied (the caller keeps its in-file defaults) and a warning is printed
    — this stops, e.g., the sim from silently loading a real-arm calibration (which would put its home
    sideways). Files without a ``convention`` tag (or no ``expected_convention``) are applied as-is."""
    import json

    if isinstance(source, dict):
        cfg = source
    else:
        with open(source) as f:
            cfg = json.load(f)
    conv = cfg.get("convention")
    if expected_convention is not None and conv is not None and conv != expected_convention:
        print(
            f"WARNING: IK calibration convention {conv!r} != active backend {expected_convention!r}; "
            "NOT applying it (home/rest/limits are convention-specific). Using in-file defaults."
        )
        return cfg
    g = globals()
    if "joint_limits_deg" in cfg:
        JOINT_LIMITS_DEG.clear()
        JOINT_LIMITS_DEG.update({k: tuple(float(x) for x in v) for k, v in cfg["joint_limits_deg"].items()})
    if "home_pose_deg" in cfg:
        READY_POSE_DEG.clear()
        READY_POSE_DEG.update({k: float(v) for k, v in cfg["home_pose_deg"].items()})
    if "rest_pose_deg" in cfg:
        RESTING_POSE_DEG.clear()
        RESTING_POSE_DEG.update({k: float(v) for k, v in cfg["rest_pose_deg"].items()})
    for key, name in (
        ("tool_length_m", "TOOL_LENGTH_M"),
        ("approach_cap_deg", "MAX_ORIENTATION_RESIDUAL_DEG"),
        ("floor_z_m", "FLOOR_Z_M"),
        ("jog_frame_yaw_deg", "JOG_FRAME_YAW_DEG"),
    ):
        if key in cfg:
            g[name] = float(cfg[key])
    for key, name in (
        ("cart_max_speed_mps", "CART_MAX_SPEED_MPS"),
        ("cart_max_accel_mps2", "CART_MAX_ACCEL_MPS2"),
        ("joint_max_speed_dps", "JOINT_MAX_SPEED_DPS"),
        ("joint_path_speed_dps", "JOINT_PATH_SPEED_DPS"),
        ("joint_path_accel_dps2", "JOINT_PATH_ACCEL_DPS2"),
        ("gripper_speed_pps", "GRIPPER_SPEED_PPS"),
    ):
        if key in cfg.get("motion", {}):
            g[name] = float(cfg["motion"][key])
    gr = cfg.get("gripper", {})
    if "servo_open_deg" in gr:
        g["GRIPPER_SERVO_OPEN_DEG"] = float(gr["servo_open_deg"])
    if "servo_closed_deg" in gr:
        g["GRIPPER_SERVO_CLOSED_DEG"] = float(gr["servo_closed_deg"])
    if "gripper_limits" in cfg:  # the gripper's calibrated open<->closed travel (0..100); clamps the command
        g["GRIPPER_LIMITS"] = tuple(float(x) for x in cfg["gripper_limits"])
    if "jog_flip_lr" in cfg:  # mirror LEFT/RIGHT for a left-handed world XY (operator frame)
        g["JOG_FLIP_LR"] = bool(cfg["jog_flip_lr"])
    if "gripper_invert" in cfg:  # servo-low = jaw open: send (100 - command) so "close" closes
        g["GRIPPER_INVERT"] = bool(cfg["gripper_invert"])
    if "workspace_radius_m" in cfg:  # table-disk radius confining the tip during manual control (None=off)
        wr = cfg["workspace_radius_m"]
        g["WORKSPACE_RADIUS_M"] = None if wr is None else float(wr)
    if "workspace_center_xy" in cfg:  # table-disk centre in the arm's OWN IK frame (real arm: ~[0, 0.381])
        g["WORKSPACE_CENTER_XY"] = tuple(float(v) for v in cfg["workspace_center_xy"])
    if "include_gripper" in cfg:
        g["INCLUDE_GRIPPER"] = bool(cfg["include_gripper"])
    return cfg


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--urdf-path", default=str(DEFAULT_URDF), help="Bundled sim URDF (with local STL meshes).")
    p.add_argument("--check", action="store_true", help="Run the headless self-check and exit.")
    p.add_argument(
        "--demo", action="store_true", help="Play a scripted trajectory in the viewer (no keyboard)."
    )
    p.add_argument(
        "--web",
        action="store_true",
        help="Browser control panel: capture keys/buttons in the browser (no terminal focus needed).",
    )
    p.add_argument(
        "--remote-port",
        type=int,
        default=None,
        help="With --web: also accept remote teleop targets over UDP on this port (see tools/latency_probe.py).",
    )
    p.add_argument(
        "--single-arm", action="store_true", help="With --web: drive a single arm (arm 1) only."
    )
    p.add_argument(
        "--gamepad",
        action="store_true",
        help="With --web: also drive the arm(s) with Joy-Con(s) (one per arm). Map first with -m tomctl.gamepad --map.",
    )
    p.add_argument(
        "--gamepad-profile",
        default=None,
        help="Joy-Con mapping JSON (from -m tomctl.gamepad --map). Default: gamepad_profile.json if present.",
    )
    p.add_argument(
        "--calibration",
        default=None,
        help="IK-control calibration JSON (from -m tomctl.calibrate.arm). The sim runs on its in-file defaults unless you pass one explicitly.",
    )
    # Remote camera streaming (the browser page publishes its sim views to LiveKit; see the panel
    # button or open with ?stream=1). Defaults are the dev/LAN key on the same room as the cameras.
    p.add_argument("--livekit-url", default="ws://localhost:7880", help="LiveKit ws URL for sim camera streaming")
    p.add_argument("--livekit-room", default="tom-cams", help="LiveKit room to publish sim cameras into")
    p.add_argument("--livekit-key", default="devkey", help="LiveKit API key (dev mode: devkey)")
    p.add_argument("--livekit-secret", default="secret", help="LiveKit API secret (dev mode: secret)")
    p.add_argument("--web-host", default="127.0.0.1",
                   help="Bind address for the web/control-panel HTTP server. Use 0.0.0.0 to reach the "
                        "sim panel (and stream its cameras) from another machine on the LAN.")
    args = p.parse_args()

    if args.gamepad:  # Linux: relaunch with a udev-enabled SDL if the bundled one can't see joysticks
        from tomctl import gamepad as _gi  # noqa: PLC0415

        _gi.ensure_udev_sdl()  # must run before pygame is imported (via the gamepad module below)

    # Calibration is opt-in for the sim: with no --calibration it runs on its in-file defaults (there
    # is no shared/legacy default file). An explicit --calibration is loaded and must exist.
    # Two guards: --check never loads (the self-check is the regression gate on the in-file defaults,
    # so it stays deterministic regardless of any file present); and a real-arm calibration is
    # refused here (expected_convention="sim") so it can't put the sim's home sideways.
    cal_path = Path(args.calibration) if args.calibration else None
    if cal_path and not cal_path.exists():
        raise SystemExit(f"Calibration file not found: {cal_path}")
    if cal_path and not args.check:
        cfg = load_calibration(cal_path, expected_convention="sim")
        if cfg.get("convention") in (None, "sim"):
            print(f"Loaded IK calibration: {cal_path}")

    urdf_path = Path(args.urdf_path).resolve()
    if not urdf_path.exists():
        raise SystemExit(
            f"Sim URDF not found: {urdf_path}\nBuild the bundle first:\n"
            "  uv run --with trimesh --with lxml --with fast-simplification "
            "python build_sim_assets.py"
        )
    model = load_urdf_model(urdf_path)

    if args.check:
        raise SystemExit(0 if run_check(model) else 1)
    if args.demo:
        run_demo(model)
    elif args.web:
        run_web(
            model,
            remote_port=args.remote_port,
            single_arm=args.single_arm,
            gamepad=args.gamepad,
            gamepad_profile=args.gamepad_profile,
            livekit_url=args.livekit_url,
            livekit_room=args.livekit_room,
            livekit_key=args.livekit_key,
            livekit_secret=args.livekit_secret,
            web_host=args.web_host,
        )
    elif args.gamepad:
        raise SystemExit("--gamepad is used together with --web (browser renderer). Add --web.")
    else:
        run_live(model)


if __name__ == "__main__":
    main()