#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
arm_meshcat_ar4.py · LateAI 站内原创教学代码
用【真实 3D 模型】（官方 STL 转出的 OBJ 网格）在 MIT Drake 中做真实模拟，
并用 Meshcat 输出可在浏览器里回放的 3D 动画。

与 arm_wave_drake_ar4.py（长方体示意模型）的区别：
  · 这里加载 ar4_real.urdf → 18 个官方零件网格，看到的就是真机外观
  · 增益不再手调：脚本先用「动能法」测出每个关节的有效惯量，再按
      kp = min(effort/0.05, 0.5·(1.5/Δt)²·I)、kd = 1.6·sqrt(kp·I)  自动整定
  · 末尾把整段动画导出成单文件 HTML（Meshcat 静态回放），可离线打开

用法:
  python3 arm_meshcat_ar4.py --html arm_real_3d_ar4.html   # 跑并导出 3D 动画（默认）
  python3 arm_meshcat_ar4.py --shape circle --hold --html none
                                                          # 跑网页画板上的圆形轨迹并实时看
  python3 arm_meshcat_ar4.py --no-meshcat --duration 20    # 只出数值与 CSV
  python3 arm_meshcat_ar4.py --q-traj q_traj.npy           # 回放网页导出的关节轨迹

三种形状同时看（各起一个进程，端口分开，互不干扰）：
  python3 arm_meshcat_ar4.py --shape circle --port 7000 --hold --html none
  python3 arm_meshcat_ar4.py --shape sine   --port 7001 --hold --html none
  python3 arm_meshcat_ar4.py --shape tri    --port 7002 --hold --html none
  然后浏览器分别打开 http://localhost:7000 / 7001 / 7002 对照

--shape {sine,circle,tri} 会用 Drake 自带逆运动学，解出与网页「3D 轨迹实验室」
同一块画板上的末端轨迹对应的关节角，再放进真实动力学里跑，并在 3D 场景里
画出「参考轨迹（蓝线）+ 末端实际轨迹（黄线）+ 当前目标点（红球）」——
球与末端之间的空隙就是真实动力学的跟踪滞后。
"""
import argparse
import math
import os

import numpy as np
from pydrake.all import (
    AddMultibodyPlantSceneGraph,
    DiagramBuilder,
    InverseKinematics,
    Meshcat,
    MeshcatVisualizer,
    MultibodyPlant,
    Parser,
    Rgba,
    RigidTransform,
    Simulator,
    Solve,
    Sphere,
)

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_URDF = os.path.join(HERE, "ar4_real.urdf")
JOINT_NAMES = ["J1", "J2", "J3", "J4", "J5", "J6"]

# 官方 xacro 限位（度）：J1 ±170, J2 -42~90, J3 -89~52, J4 ±180, J5 ±105, J6 ±180
LIMIT_DEG = [(-170, 170), (-42, 90), (-89, 52), (-180, 180), (-105, 105), (-180, 180)]

# ---------- 画板参数：与网页「AR4 真实 3D 轨迹实验室」逐字段一致 ----------
# 网页用 Y-up 显示，其坐标 (x, YC, ZP) 换到 Drake 的 Z-up 世界系即 (x, -ZP, YC)：
# 画板竖直立在正前方 y = -0.174 m、板中心高 z = 0.40 m。两处同一块板，
# 所以网页上画出来的形状，和这里的 --shape 是同一个目标点集。
BOARD_ZP, BOARD_YC, BOARD_R = 0.174, 0.40, 0.055   # 前伸 / 中心高度 / 圆半径
SINE_L, SINE_A, SINE_W = 0.075, 0.045, 62.0        # 正弦：半长 / 幅值 / 角频率
SHAPE_PERIOD = 5.0                                  # 每圈（每程）5 s，与网页一致
# IK 起点（与网页相同）：J4、J6 全程锁 0，只让 J1/J2/J3/J5 去解位置
Q_SEED = np.array([0.0, -0.5, 0.7, 0.0, 0.2, 0.0])

# ---------- 轨迹加粗 ----------
# 坑：Meshcat 前端渲染线用的是 THREE.LineBasicMaterial，线宽最终交给 gl.lineWidth，
# 而 WebGL 在 macOS/Chrome 下把它硬限制成 1 像素 —— 所以 SetLine 的 line_width
# 不管写 3 还是 300，浏览器里永远是一条 1 px 细线，远看就糊成一团。
# 对策：轨迹额外用「一串密排的小球」铺出来，球够密就连成一条真正有粗细的线
#（原来的细线保留，用来表示精确路径）。网页版 arm_traj_3d.html 同理，见那边的点串。
REF_BEAD_R, REF_BEAD_STRIDE = 0.0035, 2       # 参考轨迹(蓝)：球半径 m / 每几点一颗
TRAIL_BEAD_R, TRAIL_BEAD_STRIDE = 0.0055, 4   # 实际轨迹(黄)：0.02 s 一点 ≈ 8 mm 一颗
TRAIL_BEAD_SLOTS = 750                        # 球池槽位数（懒创建 + 环形复用，覆盖满 3000 点）
REF_RGBA = Rgba(0.42, 0.55, 1.0, 0.85)        # 参考轨迹：蓝
TRAIL_RGBA = Rgba(1.0, 0.84, 0.3, 0.95)       # 实际轨迹：黄


def shape_points(shape, n=240):
    """画板上的末端目标点，返回 n×3（Drake 世界系）。

    三种形状与网页 arm_traj_3d.html 的 targetAt() 完全一致：
      sine   左右往返一次（去程 -L→+L、回程 +L→-L），板面留下正弦曲线
      circle 逆时针一整圈
      tri    从正上方顶点起逆时针一圈（三条边等分一周）
    三者都是闭合周期路径，首尾相接，可直接循环播放。
    """
    ts = np.linspace(0.0, 1.0, n, endpoint=False)
    pts = np.zeros((n, 3))
    if shape == "circle":
        a = 2.0 * np.pi * ts
        pts[:, 0] = BOARD_R * np.cos(a)
        pts[:, 2] = BOARD_YC + BOARD_R * np.sin(a)
    elif shape == "tri":
        t3 = ts * 3.0
        i3 = np.floor(t3).astype(int) % 3
        tt = t3 - np.floor(t3)
        a0 = np.pi / 2.0 + 2.0 * np.pi * i3 / 3.0
        a1 = np.pi / 2.0 + 2.0 * np.pi * ((i3 + 1) % 3) / 3.0
        x0, z0 = BOARD_R * np.cos(a0), BOARD_YC + BOARD_R * np.sin(a0)
        x1, z1 = BOARD_R * np.cos(a1), BOARD_YC + BOARD_R * np.sin(a1)
        pts[:, 0] = x0 + (x1 - x0) * tt
        pts[:, 2] = z0 + (z1 - z0) * tt
    else:                                        # sine
        left = ts < 0.5
        u = np.where(left, -SINE_L + 4.0 * SINE_L * ts,
                     SINE_L - 4.0 * SINE_L * (ts - 0.5))
        pts[:, 0] = u
        pts[:, 2] = BOARD_YC + SINE_A * np.sin(SINE_W * u)
    pts[:, 1] = -BOARD_ZP                       # 全部落在画板平面内
    return pts


def solve_ik_path(plant, frame, pts, seed=None):
    """沿画板路径逐点解 IK，返回 N×6 关节角。

    与网页一致：J4、J6 锁 0（末端保持固定姿态），只解 J1/J2/J3/J5。

    关键在于「逐点续解」：每点以上一点的解作初值，再加一个很轻的正则项把
    解往初值附近拉。若像网页那样只用固定的初始猜测、或稀疏采样，
    Drake 的 IK 会中途跳到另一个分支 —— 实测 J5 从 0.19 突跳到 1.62 rad，
    关节轨迹直接断裂，真机根本执行不了。密集采样 + 续解后，
    实测步间最大变化 0.008 rad、三轨迹末端跟踪误差均 ≤ 0.35 mm。
    """
    q = (Q_SEED if seed is None else np.asarray(seed, dtype=float)).copy()
    nq = plant.num_positions()
    solve_idx = [0, 1, 2, 4]                     # J1/J2/J3/J5
    W = np.zeros((len(solve_idx), nq))
    for r, i in enumerate(solve_idx):
        W[r, i] = 1.0
    Q = np.zeros((len(pts), nq))
    nfail = 0
    for k, p in enumerate(pts):
        ik = InverseKinematics(plant)
        qv = ik.q()
        prog = ik.prog()
        ik.AddPositionConstraint(frame, [0, 0, 0], plant.world_frame(),
                                 p - 2e-4, p + 2e-4)
        prog.AddQuadraticErrorCost(1e-3 * (W.T @ W), q, qv)
        prog.AddBoundingBoxConstraint(q[3], q[3], qv[3:4])   # J4 锁 0
        prog.AddBoundingBoxConstraint(q[5], q[5], qv[5:6])   # J6 锁 0
        prog.SetInitialGuess(qv, q)
        res = Solve(prog)
        if res.is_success():
            q = np.array(res.GetSolution())
        else:
            nfail += 1
        Q[k] = q
    if nfail:
        print("[ar4_real] 警告：IK 有 %d/%d 点未收敛" % (nfail, len(pts)))
    return Q


def cmd_angle(i, t, freq=0.5, amp_deg=25.0):
    """第 i 个关节的目标角度(rad)：以限位区间中点为基准做正弦摆动。

    AR4 的 J2/J3 限位不对称，不能以 0 为中心大摆，否则会在真实动力学里撞限位。
    """
    lo, hi = np.radians(LIMIT_DEG[i])
    mid = 0.5 * (lo + hi)
    span = 0.5 * (hi - lo)
    amp = min(math.radians(amp_deg), 0.6 * span)
    return mid + amp * math.sin(freq * t + i * math.pi / 3.0)


def cmd_angle_dot(i, t, freq=0.5, amp_deg=25.0):
    lo, hi = np.radians(LIMIT_DEG[i])
    span = 0.5 * (hi - lo)
    amp = min(math.radians(amp_deg), 0.6 * span)
    return amp * freq * math.cos(freq * t + i * math.pi / 3.0)


def effective_inertia(plant, ctx, vel_idx):
    """动能法测各关节有效转动惯量：令 v = e_i，则 T = 0.5·I_i ⇒ I_i = 2T。"""
    g_backup = plant.gravity_field().gravity_vector().copy()
    plant.mutable_gravity_field().set_gravity_vector([0, 0, 0])
    out = []
    for s in vel_idx:
        v = np.zeros(plant.num_velocities())
        v[s] = 1.0
        plant.SetVelocities(ctx, v)
        out.append(2.0 * plant.CalcKineticEnergy(ctx))
    plant.SetVelocities(ctx, np.zeros(plant.num_velocities()))
    plant.mutable_gravity_field().set_gravity_vector(g_backup)
    return np.array(out)


def auto_gain(effort, inertia, dt):
    """按扭矩上限(线性域)与离散化稳定判据自动整定 PD。

    kd 额外受「微分时间常数 τ = I/kd ≥ 2Δt」约束：腕部轴惯量比 J2 小
    2~3 个量级（J6 仅 9e-5 kg·m²），不约束时 τ 小于控制周期，一步内
    速度反馈即把力矩顶到上限，形成高频极限环（实测 J6 角速度飙到
    ±47 rad/s，而目标仅 0.22 rad/s）。
    """
    I = np.array(inertia)
    kp = np.minimum(np.array(effort) / 0.05, 0.5 * (1.5 / dt) ** 2 * I)
    kd = np.minimum(1.6 * np.sqrt(kp * I), I / (2.0 * dt))
    return kp, kd


def strip_collisions(urdf_path):
    """把 URDF 里的 <collision> 块剥掉，返回可加载的文件路径。

    为什么需要：官方 URDF 给每个连杆只配了一个粗略的长方体 collision。
    相邻/邻近连杆的这些盒子互相穿插（例如 link4 的盒子覆盖 z∈[-0.22, 0]，
    link5 的盒子又落在同一区间），Drake 的接触求解器于是施加巨大的穿透
    修正力，把腕部关节"锁死"：实测 J4 饱和占比 89%、跟踪误差 0.43 rad，
    而剥离碰撞后立刻降到 0% / 0.007 rad。
    教学仿真关心的是关节控制与运动学，真实碰撞模型应另行建立。
    注意：生成的临时文件必须与源 URDF 同目录，否则内部相对 mesh 路径会失效。
    """
    import re
    with open(urdf_path, encoding="utf-8") as f:
        txt = f.read()
    if "<collision>" not in txt:
        return urdf_path
    txt = re.sub(r"[ \t]*<collision>.*?</collision>\s*\n?", "", txt, flags=re.S)
    out = os.path.join(os.path.dirname(os.path.abspath(urdf_path)),
                       "_nocol_" + os.path.basename(urdf_path))
    with open(out, "w", encoding="utf-8") as f:
        f.write(txt)
    print("[ar4_real] 已剥离 collision -> %s（源文件不动）" % os.path.basename(out))
    return out


def load_traj(path):
    """读取网页导出的 q_traj.npy（N×6 关节角，rad）。"""
    q = np.load(path)
    q = np.atleast_2d(q)
    if q.shape[1] != 6 and q.shape[0] == 6:
        q = q.T
    if q.shape[1] != 6:
        raise ValueError("q_traj 形状应为 N×6，实际为 %s" % (q.shape,))
    return q


def main():
    ap = argparse.ArgumentParser(description="Annin AR4 (MK5) 真实 3D 模型 · Drake + Meshcat 模拟")
    ap.add_argument("--urdf", default=DEFAULT_URDF, help="真实 URDF 路径（含 OBJ 网格）")
    ap.add_argument("--duration", type=float, default=12.0, help="仿真时长(秒)")
    ap.add_argument("--dt", type=float, default=0.001, help="控制刷新周期(秒，1kHz)")
    ap.add_argument("--time-step", type=float, default=0.001,
                    help="plant 积分步长；默认 0.001（离散半隐欧拉，与 arm_wave 一致）。"
                         "传 0 会切成连续积分器 —— 本模型腕部惯量极小、系统刚性很强，"
                         "连续积分器实测 4 分钟跑不完 1 秒仿真，不建议")
    ap.add_argument("--keep-collisions", action="store_true",
                    help="保留 URDF 自带的碰撞体；默认剥离（原因见 strip_collisions）")
    ap.add_argument("--hold", action="store_true",
                    help="跑完后保持 3D 服务常驻，浏览器实时查看（配合 --html none 跳过静态导出）")
    ap.add_argument("--port", type=int, default=7000,
                    help="Meshcat 网页端口，默认 7000。想同时看多种形状时用不同端口各起一个："
                         "例如 --shape circle --port 7000、--shape sine --port 7001、"
                         "--shape tri --port 7002，浏览器分别打开对应地址即可对照")
    ap.add_argument("--freq", type=float, default=0.5, help="摆动圆频率(rad/s)")
    ap.add_argument("--amp-deg", type=float, default=25.0, help="摆动幅值(°)")
    ap.add_argument("--q-traj", default=None, help="回放网页导出的 q_traj.npy（N×6）")
    ap.add_argument("--traj-period", type=float, default=5.0,
                    help="q_traj 对应的原始周期(秒)，用于线性重采样")
    ap.add_argument("--shape", choices=["wave", "sine", "circle", "tri"], default="wave",
                    help="wave=各关节正弦摆动（默认）；sine/circle/tri=网页画板上的"
                         "正弦曲线/圆形/三角形：用 Drake 逆运动学解出关节角，再跑真实动力学")
    ap.add_argument("--shape-points", type=int, default=240,
                    help="画板轨迹的 IK 采样点数（每圈）；越大越平滑，IK 也越慢")
    ap.add_argument("--save-q", default=None,
                    help="把解出的关节轨迹存成 .npy，之后可用 --q-traj 回放")
    ap.add_argument("--no-trail", action="store_true", help="不在 3D 场景里画末端运动轨迹")
    ap.add_argument("--publish-period", type=float, default=0.02, help="3D 可视化刷新周期(秒)")
    ap.add_argument("--csv", default=os.path.join(HERE, "ar4_real_trail.csv"))
    ap.add_argument("--html", default=os.path.join(HERE, "arm_real_3d_ar4.html"),
                    help="导出的 meshcat 3D 动画 HTML")
    ap.add_argument("--no-meshcat", action="store_true", help="不启动 3D，只出数值")
    args = ap.parse_args()

    use_vis = not args.no_meshcat
    if use_vis:
        builder = DiagramBuilder()
        plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=args.time_step)
    else:
        plant = MultibodyPlant(time_step=args.time_step)

    Parser(plant).AddModels(args.urdf if args.keep_collisions
                            else strip_collisions(args.urdf))
    # 真实 URDF 里 base 是根连杆（除非带 world_joint）：必须焊到世界，否则整臂自由落体
    try:
        plant.GetJointByName("world_joint")
    except RuntimeError:
        plant.WeldFrames(plant.world_frame(), plant.GetBodyByName("base").body_frame())
    # URDF 不自动创建执行器：revolute 关节默认被动，必须显式 AddJointActuator，
    # 且必须在 plant.Finalize() 之前完成（Finalize 后模型就锁定了）。
    # 扭矩上限按 AR4 真实配置估算：J1~J3 为 NEMA23 + 减速箱，J4~J6 为 NEMA17 + 减速箱。
    eff_guess = [20.0, 20.0, 15.0, 8.0, 8.0, 5.0]
    for nm, ef in zip(JOINT_NAMES, eff_guess):
        plant.AddJointActuator(nm, plant.GetJointByName(nm), effort_limit=ef)
    plant.Finalize()

    n = len(JOINT_NAMES)
    joints = [plant.GetJointByName(nm) for nm in JOINT_NAMES]

    if use_vis:
        meshcat = Meshcat(port=args.port)   # 指定端口，方便把链接写进教程页/卡片
        MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)
        diagram = builder.Build()
        sim = Simulator(diagram)
        root_ctx = sim.get_mutable_context()
        ctx = plant.GetMyContextFromRoot(root_ctx)
    else:
        sim = Simulator(plant)
        ctx = sim.get_mutable_context()
        root_ctx = ctx

    acts = [plant.GetJointActuatorByName(nm) for nm in JOINT_NAMES]
    effort = np.array([a.effort_limit() for a in acts])
    vel_idx = [j.velocity_start() for j in joints]
    tool = plant.GetBodyByName("tool_link")

    inertia = effective_inertia(plant, ctx, vel_idx)
    kp_arr, kd_arr = auto_gain(effort, inertia, args.dt)
    print("[ar4_real] joints=%d · effort=%s" % (n, np.array2string(effort, precision=1)))
    print("[ar4_real] 有效惯量=%s" % np.array2string(inertia, precision=5))
    print("[ar4_real] 自动整定 kp=%s" % np.array2string(kp_arr, precision=1))
    print("[ar4_real] 自动整定 kd=%s" % np.array2string(kd_arr, precision=2))

    use_shape = args.shape != "wave"
    ref_pts = None
    Q = None
    seg_dt = 0.0
    if args.q_traj:
        Q = load_traj(args.q_traj)
        seg_dt = args.traj_period / max(len(Q) - 1, 1)
        print("[ar4_real] 回放 q_traj: %d 帧 → 重采样到 %.1f s（原周期 %.1f s）"
              % (len(Q), args.duration, args.traj_period))
    elif use_shape:
        ref_pts = shape_points(args.shape, args.shape_points)
        print("[ar4_real] 画板轨迹 %s：%d 个目标点，解 IK 中（J1/J2/J3/J5 求解，J4/J6 锁 0）…"
              % (args.shape, len(ref_pts)))
        Q = solve_ik_path(plant, tool.body_frame(), ref_pts)
        seg_dt = SHAPE_PERIOD / len(Q)
        ctx_fk = plant.CreateDefaultContext()
        e = []
        for p, qq in zip(ref_pts, Q):
            plant.SetPositions(ctx_fk, qq)
            e.append(np.linalg.norm(
                plant.EvalBodyPoseInWorld(ctx_fk, tool).translation() - p))
        print("[ar4_real] IK 完成 · 每圈 %.1f s · 末端几何误差 max=%.2f mm · "
              "J1[%.2f, %.2f] J2[%.2f, %.2f] J3[%.2f, %.2f] J5[%.2f, %.2f] rad"
              % (SHAPE_PERIOD, np.max(e) * 1000,
                 Q[:, 0].min(), Q[:, 0].max(), Q[:, 1].min(), Q[:, 1].max(),
                 Q[:, 2].min(), Q[:, 2].max(), Q[:, 4].min(), Q[:, 4].max()))
        if args.save_q:
            np.save(args.save_q, Q)
            print("[ar4_real] 关节轨迹已导出 -> %s（可喂给 arm_wave_drake_ar4.py --q-traj）"
                  % args.save_q)

    u_port = plant.get_actuation_input_port()
    u = np.zeros(u_port.size())
    u_port.FixValue(ctx, u)      # 端口属于 plant，必须用 plant 的 context

    if Q is not None and use_shape:
        # 从轨迹起点起步，免得开机瞬间从零位猛甩到画板起点
        plant.SetPositions(ctx, Q[0])

    sim.Initialize()
    if use_vis:
        meshcat.StartRecording()
        if ref_pts is not None:
            # 画板：竖在正前方 y=-0.174 m、中心高 0.40 m，尺寸与网页画板一致。
            # 蓝线 = 参考轨迹，黄线 = 末端实际轨迹，红球 = 当前目标点；
            # 红球与末端之间的空隙就是真实动力学的跟踪滞后。
            hw, hh = 0.15, 0.11
            c = np.array([0.0, -BOARD_ZP, BOARD_YC])
            corners = c + np.array([[-hw, 0.0, -hh], [hw, 0.0, -hh],
                                    [hw, 0.0, hh], [-hw, 0.0, hh]])
            meshcat.SetLineSegments("/board/frame", np.asfortranarray(corners.T),
                                    np.asfortranarray(np.roll(corners, -1, axis=0).T),
                                    line_width=2.0, rgba=Rgba(0.35, 0.45, 0.9, 0.5))
            meshcat.SetLine("/board/ref", np.asfortranarray(ref_pts.T), line_width=3.0,
                            rgba=REF_RGBA)
            # 参考轨迹铺粗：密排蓝球，球心间距 ≈2.8 mm、直径 7 mm，连起来即一条粗线
            for k in range(0, len(ref_pts), REF_BEAD_STRIDE):
                meshcat.SetObject("/board/ref_bead/%d" % k, Sphere(REF_BEAD_R),
                                  rgba=REF_RGBA)
                meshcat.SetTransform("/board/ref_bead/%d" % k, RigidTransform(ref_pts[k]))
            meshcat.SetObject("/board/target", Sphere(0.006),
                              rgba=Rgba(1.0, 0.36, 0.45, 1.0))
    t = 0.0
    next_print = 0.0
    next_trail = 0.0
    trail = []
    trail_n = 0          # 已采样点数：决定球串槽位
    trail_slots = set()  # 已创建过的球槽位（懒创建，短仿真不会白建一堆空球）
    err_max = np.zeros(n)
    sat_steps = np.zeros(n)
    rows = []
    while t < args.duration - 1e-12:
        q = np.array([joints[i].get_angle(ctx) for i in range(n)])
        qd = plant.GetVelocities(ctx)[vel_idx]
        ref_now = None
        if Q is not None:
            if args.q_traj:                    # 网页导出的有限帧序列：两端夹取
                x = np.clip(t / args.traj_period * (len(Q) - 1), 0, len(Q) - 1)
                i0 = int(np.floor(x))
                i1 = min(i0 + 1, len(Q) - 1)
                f = x - i0
            else:                              # 画板轨迹是闭合周期：首尾环绕
                x = (t / SHAPE_PERIOD % 1.0) * len(Q)
                i0 = int(x) % len(Q)
                i1 = (i0 + 1) % len(Q)
                f = x - math.floor(x)
                ref_now = ref_pts[i0] * (1.0 - f) + ref_pts[i1] * f
            cmd = Q[i0] * (1.0 - f) + Q[i1] * f
            cmd_dot = (Q[i1] - Q[i0]) / seg_dt
        else:
            cmd = np.array([cmd_angle(i, t, args.freq, args.amp_deg) for i in range(n)])
            cmd_dot = np.array([cmd_angle_dot(i, t, args.freq, args.amp_deg)
                                for i in range(n)])
        tau_g = plant.CalcGravityGeneralizedForces(ctx)[vel_idx]
        u[:] = kp_arr * (cmd - q) + kd_arr * (cmd_dot - qd) - tau_g
        if not np.isfinite(u).all():
            print("NaN! t=%.4f\n q=%s\n qd=%s\n cmd=%s\n tau_g=%s" % (t, q, qd, cmd, tau_g))
            break
        u[:] = np.clip(u, -effort, effort)
        sat_steps += (np.abs(u) >= effort - 1e-9).astype(float)
        u_port.FixValue(ctx, u)
        sim.AdvanceTo(t + args.dt)
        t += args.dt

        if t >= next_trail:
            # 末端实际轨迹：定期采样 tool_link 的世界坐标，连成一条黄线画在场景里。
            # 这条线就是「真机跑出来的路径」，可与蓝线（参考轨迹）逐点对照。
            next_trail += args.publish_period
            xyz_v = plant.EvalBodyPoseInWorld(ctx, tool).translation()
            trail.append(np.array(xyz_v))
            if len(trail) > 3000:
                del trail[0]
            if use_vis:
                meshcat.SetSimulationTime(t)
                if not args.no_trail:
                    meshcat.SetLine("/trail", np.asfortranarray(np.array(trail).T),
                                    line_width=3.0, rgba=TRAIL_RGBA)
                    # 实际轨迹也铺粗：每 TRAIL_BEAD_STRIDE 个采样点放一颗黄球，
                    # 槽位环形复用 —— 与 trail 本身 3000 点的滑动窗口语义一致。
                    if trail_n % TRAIL_BEAD_STRIDE == 0:
                        slot = (trail_n // TRAIL_BEAD_STRIDE) % TRAIL_BEAD_SLOTS
                        bead_path = "/trail_bead/%d" % slot
                        if slot not in trail_slots:
                            meshcat.SetObject(bead_path, Sphere(TRAIL_BEAD_R),
                                              rgba=TRAIL_RGBA)
                            trail_slots.add(slot)
                        meshcat.SetTransform(bead_path, RigidTransform(xyz_v))
                if ref_now is not None:
                    meshcat.SetTransform("/board/target", RigidTransform(ref_now))
            trail_n += 1

        if t >= next_print:
            next_print += 0.5
            if t >= 2.0:
                err_max = np.maximum(err_max, np.abs(cmd - q))
            xyz = plant.EvalBodyPoseInWorld(ctx, tool).translation()
            rows.append([t] + list(q) + list(xyz))
            if abs(t - round(t)) < 1e-9:
                print("t=%.1fs 实测(rad): %s  末端XYZ(m): %.3f %.3f %.3f"
                      % (t, " ".join("%8.3f" % v for v in q), *xyz))

    if use_vis:
        meshcat.StopRecording()
        meshcat.PublishRecording()
        print("[ar4_real] 3D 服务已就绪 -> %s" % meshcat.web_url())
        if args.html and args.html.lower() != "none":
            html = meshcat.StaticHtml()
            with open(args.html, "w") as f:
                f.write(html)
            print("[ar4_real] 3D 动画已导出 -> %s (%.1f MB)" % (args.html, len(html) / 1e6))

    with open(args.csv, "w") as f:
        f.write("t," + ",".join("q%d" % (i + 1) for i in range(n)) + ",x,y,z\n")
        for r in rows:
            f.write(",".join("%.6f" % v for v in r) + "\n")
    n_steps = max(int(args.duration / args.dt), 1)
    print("[summary] 最大跟踪误差(rad): %s" % " ".join("%.3f" % v for v in err_max))
    print("[summary] 扭矩饱和占比(%%):  %s"
          % " ".join("%.0f" % (100 * s / n_steps) for s in sat_steps))
    if rows:
        xyz_a = np.array([r[1 + n:] for r in rows])
        print("[summary] 末端轨迹范围(m): x[%.3f, %.3f] y[%.3f, %.3f] z[%.3f, %.3f]"
              % (xyz_a[:, 0].min(), xyz_a[:, 0].max(), xyz_a[:, 1].min(),
                 xyz_a[:, 1].max(), xyz_a[:, 2].min(), xyz_a[:, 2].max()))
        if ref_pts is not None and trail:
            # 末端实际路径 vs 画板参考路径：差值 = 控制滞后 + 几何误差，即「真机代价」
            tp = np.array(trail)
            d = np.array([np.min(np.linalg.norm(ref_pts - p, axis=1)) for p in tp])
            print("[summary] 末端偏离参考路径(m): 最大=%.4f 平均=%.4f"
                  % (d.max(), d.mean()))
    print("[done] 轨迹已保存 -> %s" % args.csv)

    if use_vis and args.hold:
        # 必须放在最后：保持进程存活，meshcat 服务才不会随脚本退出而关闭。
        # 动画已 PublishRecording，浏览器底部播放条可反复回放。
        import time
        print("[ar4_real] 常驻中，浏览器打开 %s 查看；Ctrl+C 退出" % meshcat.web_url())
        try:
            while True:
                time.sleep(0.5)
        except KeyboardInterrupt:
            print("[ar4_real] 已退出")


if __name__ == "__main__":
    main()