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

与 arm_wave_drake_faze4.py（长方体示意模型）的区别：
  · 这里加载 faze4_real.urdf → 7 个官方零件网格，看到的就是真机外观
  · 增益不手调：先用「动能法」测各关节有效惯量，再按
      kp = min(effort/0.05, 0.4·I/Δt²)、kd = min(1.6·sqrt(kp·I), I/(2Δt))
    自动整定（Faze4 腕部惯量极小，kp 上限必须卡离散稳定判据，见下）
  · 末尾把整段动画导出成单文件 HTML（Meshcat 静态回放），可离线打开

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

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

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

⚠️ Faze4 的「惯量鸿沟」：腕部 J6 只带动约 20 g 夹爪，关节侧有效惯量约
   6.4e-6 kg·m²；而 J2 带着整条大臂，约 0.4 kg·m²，相差四五个量级。
   离散 PD 的无量纲刚度 b = Δt²·kp/I 一旦越过约 0.5 就会出现步频极限环，
   故本脚本默认 Δt = 0.2 ms（5 kHz），并对 kp 取 0.4·I/Δt² 的上限
   （与 arm_wave_drake_faze4.py 同一套判据，详见那边文件顶部的推导）。
   真实 Faze4 固件也跑在数 kHz 控制环上，所以 5 kHz 既贴近真机又数值稳定。
"""
import argparse
import math
import os

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

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_URDF = os.path.join(HERE, "faze4_real.urdf")

# 官方 URDF 的关节名是克罗地亚语（Faze4 项目来自克罗地亚），依次对应 J1~J6：
#   回转底座 / 大臂 / 肘 / 前臂滚转 / 腕俯仰 / 末端滚转
JOINT_NAMES = ["rotary_base", "nadlaktica", "lakat", "podlaktica", "saka", "hvataljka"]
JOINT_LABELS = ["J1", "J2", "J3", "J4", "J5", "J6"]

# 关节限位（度）：官方六个关节均为 continuous，这里按机械结构取（与 URDF 一致）
LIMIT_DEG = [(-180, 180), (-120, 120), (-150, 150), (-180, 180), (-120, 120), (-180, 180)]

# 扭矩上限（N·m）：六个关节均为 NEMA17 步进电机 + 内置摆线针轮减速箱，
# 由「保持扭矩 × 减速比 × 传动效率」估算（真机标定后请替换）。
EFFORT = [7.9, 8.5, 4.7, 3.7, 3.5, 6.0]

# ---------- 画板参数：与网页「Faze4 真实 3D 轨迹实验室」逐字段一致 ----------
# 网页用 Y-up 显示，其坐标 (x, YC, ZP) 换到 Drake 的 Z-up 世界系即 (x, -ZP, YC)：
# 画板竖直立在正前方 y = -0.48 m、板横向中心 x = 0.08 m（对准 J1 轴）、板中心高 0.55 m。
# 两处是同一块板，所以网页上画出来的形状与这里的 --shape 是同一个目标点集。
BOARD_ZP, BOARD_XC, BOARD_YC, BOARD_R = 0.48, 0.08, 0.55, 0.16  # 前伸/横向中心/高度/圆半径
SINE_L, SINE_A = 0.19, 0.12                                     # 正弦：半长 / 幅值
SINE_W = math.pi / 0.19                                         # 角频率（= π/L，与网页同）
SHAPE_PERIOD = 5.0                                               # 每圈（每程）5 s，与网页一致
# IK 起点（与网页同）：J4、J6 全程锁 0，只让 J1/J2/J3/J5 去解位置
Q_SEED = np.array([0.0104, 0.1795, -0.5049, 0.0, -0.0938, 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.0045, 2       # 参考轨迹(蓝)：球半径 m / 每几点一颗
TRAIL_BEAD_R, TRAIL_BEAD_STRIDE = 0.0055, 2   # 实际轨迹(黄)：0.02 s 一点 ≈ 4.6 mm 一颗
TRAIL_BEAD_SLOTS = 900                        # 球池槽位数（懒创建 + 环形复用）
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 世界系 Z-up）。

    三种形状与网页 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_XC + 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] = BOARD_XC + 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] = BOARD_XC + 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 会中途
    跳到另一个分支，关节轨迹直接断裂，真机根本执行不了。密集采样 + 续解后，
    关节轨迹在本机实测步间最大变化约 0.01 rad 量级。
    """
    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("[faze4_real] 警告：IK 有 %d/%d 点未收敛" % (nfail, len(pts)))
    return Q


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

    Faze4 六个关节的限位都关于 0 对称，故零位即 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)   # 留 40% 余量，避免顶到限位
    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):
    """目标角速度(rad/s)，即上式的解析导数（用于速度前馈，减小跟踪相位滞后）。"""
    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。

    kp = min(effort/0.05,  KP_B_MAX · I / Δt²)
      前项保证 0.05 rad 误差内电机不饱和（线性工作区）；
      后项是离散时间 PD 的稳定上限：无量纲刚度 b = Δt²·kp/I 必须留裕量
      （Faze4 实测 b ≳ 0.5 即出现步频极限环），取 KP_B_MAX = 0.4。
    kd = min(1.6·sqrt(kp·I),  I/(2Δt))
      前项是临界阻尼略偏大的系数；后项保证微分时间常数 τ = I/kd ≥ 2Δt。

    对 Faze4 的腕部轴，I/(2Δt) 总是更小的那个，故 kd 恒等于 I/(2Δt)；
    真正决定「能不能跑」的是 kp：J2 的 I 是 J6 的数万倍，同一 Δt 下 J6 的
    kp 上限被压得极低 —— 这正是 Faze4 必须把 --dt 取到 0.2 ms 的根本原因。
    """
    I = np.array(inertia)
    KP_B_MAX = 0.4   # 无量纲刚度 b = Δt²·kp/I 的上限（含 20% 稳定裕量）
    kp = np.minimum(np.array(effort) / 0.05, KP_B_MAX * I / dt ** 2)
    kd = np.minimum(1.6 * np.sqrt(kp * I), I / (2.0 * dt))
    return kp, kd


def disable_self_collision(plant, scene_graph):
    """排除同一模型实例内部的自碰撞（Faze4 必需，不是可选优化）。

    为什么需要：Faze4 的 URDF 碰撞体是「3D 打印件实体」的粗略近似，相邻连杆
    在关节处本来就是互相咬合的（减速箱/轴承座嵌进对方腔体里），所以只要接了
    SceneGraph 就会算出持续的自碰撞接触力 —— 实测初始位姿下就有接触点，把
    J4 顶到 90% 以上饱和并颤振。无头分支没有 SceneGraph、压根不做接触计算，
    所以看不到这个现象（这正是两个分支结果不一致的原因）。
    本教程场景里没有地面也没有被抓物，自碰撞纯属网格咬合造成的伪力，
    用官方推荐的做法整体排除即可（真机做碰撞检测时才会保留它）。
    """
    mi = plant.GetBodyByName("tool_link").model_instance()
    geom_set = GeometrySet()
    for b_idx in plant.GetBodyIndices(mi):
        for g_id in plant.GetCollisionGeometriesForBody(plant.get_body(b_idx)):
            geom_set.Add(g_id)
    decl = CollisionFilterDeclaration()
    decl.ExcludeWithin(geom_set)
    scene_graph.collision_filter_manager().Apply(decl)


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="Faze4 六轴机械臂 真实 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.0002,
                    help="控制刷新周期(秒)，默认 0.2 ms / 5 kHz —— Faze4 腕部惯量极小，"
                         "1 kHz 下 kp 上限会被压到无法跟轨迹，见文件顶部说明")
    ap.add_argument("--time-step", type=float, default=None,
                    help="plant 积分步长；默认跟随 --dt（离散半隐欧拉，与真机固件的"
                         "定周期控制环同构）。传 0 会切成连续积分器：本机实测也能跑通"
                         "（1 s 仿真约 2.4 s 墙钟），但控制输入仍按 --dt 采样保持，"
                         "既不比离散快、也失去了与固件一致的可复现性，故不推荐")
    ap.add_argument("--self-collisions", action="store_true",
                    help="保留模型内部自碰撞；默认排除（原因见 disable_self_collision）")
    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, "faze4_real_trail.csv"))
    ap.add_argument("--html", default=os.path.join(HERE, "arm_real_3d_faze4.html"),
                    help="导出的 meshcat 3D 动画 HTML")
    # Meshcat 默认相机离场景很远，机械臂在 1400×900 画面里只有几十像素高，
    # 导出的单文件动画几乎看不出在动。这里默认给一个右前方俯视的近景视角。
    ap.add_argument("--camera", default="1.35,-1.45,1.00",
                    help="相机位置 x,y,z（米），默认右前方上方俯视整条臂")
    ap.add_argument("--look-at", default="0.05,-0.24,0.42",
                    help="相机注视点 x,y,z（米），默认取机械臂与画板之间")
    ap.add_argument("--no-meshcat", action="store_true", help="不启动 3D，只出数值")
    args = ap.parse_args()

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

    # 直接加载官方 URDF：惯量保持原样（不做转子惯量折算），
    # 由「缩小控制周期」而非「篡改惯量」来保证腕部数值稳定。
    Parser(plant).AddModels(args.urdf)

    # faze4_real.urdf 自带 <link name="world"> + world_joint，Drake 直接把它当世界刚体；
    # 若换成不含 world 的 URDF（如官方原始 xacro 展开件），必须手动把 base 焊到世界，
    # 否则整臂会自由落体。注意：不能用 GetBodyByName("world") 判断 —— Drake 的 world
    # body 恒名为 "world"，永远查得到，会漏掉焊接这一步。
    try:
        plant.GetJointByName("world_joint")
    except RuntimeError:
        plant.WeldFrames(plant.world_frame(), plant.GetBodyByName("base_link").body_frame())

    # URDF 不自动创建执行器（revolute 关节默认被动），需显式 AddJointActuator
    # 并声明扭矩上限 effort（须在 Finalize 之前），关节才变成「可控关节」。
    for nm, ef in zip(JOINT_NAMES, EFFORT):
        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:
        if not args.self_collisions:
            disable_self_collision(plant, scene_graph)
        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("[faze4_real] joints=%d · 控制周期=%.4f s · effort=%s"
          % (n, args.dt, np.array2string(effort, precision=1)))
    print("[faze4_real] 有效惯量=%s" % np.array2string(inertia, precision=6))
    print("[faze4_real] 自动整定 kp=%s" % np.array2string(kp_arr, precision=1))
    print("[faze4_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("[faze4_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("[faze4_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("[faze4_real] IK 完成 · 每圈 %.1f s · 末端几何误差 max=%.2f mm · "
              "J1[%.3f, %.3f] J2[%.3f, %.3f] J3[%.3f, %.3f] J5[%.3f, %.3f] 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("[faze4_real] 关节轨迹已导出 -> %s（可喂给 arm_wave_drake_faze4.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

    # 起始位姿直接落在命令轨迹的起点：Faze4 腕部惯量极小（J6 仅 6.4e-6 kg·m²），
    # 若从全零姿态起步，第一个控制周期就会吃到满幅误差 —— 腕部瞬间发散，
    # J2 也会在启动瞬态顶到扭矩上限。故无论哪种模式都从命令起点起步
    #（与 arm_wave_drake_faze4.py 同一处理）。
    q_ini = (np.array(Q[0], dtype=float) if Q is not None
             else np.array([cmd_angle(i, 0.0, args.freq, args.amp_deg) for i in range(n)]))
    plant.SetPositions(ctx, q_ini)
    plant.SetVelocities(ctx, np.zeros(plant.num_velocities()))

    sim.Initialize()
    if use_vis:
        # 必须在 StartRecording 之前定好相机：否则静态 HTML 会沿用 Meshcat 默认
        # 的远距离全景视角，机械臂小到看不出运动（这是本脚本导出动画的默认坑）。
        cam = np.array([float(v) for v in args.camera.split(",")])
        tgt = np.array([float(v) for v in args.look_at.split(",")])
        meshcat.SetCameraPose(cam, tgt)
        meshcat.StartRecording()
        if ref_pts is not None:
            # 画板：竖在正前方 y=-0.48 m、横向中心 x=0.08 m、中心高 0.55 m，
            # 板面 0.50 × 0.44 m，与网页画板一致。
            # 蓝线 = 参考轨迹，黄线 = 末端实际轨迹，红球 = 当前目标点；
            # 红球与末端之间的空隙就是真实动力学的跟踪滞后。
            hw, hh = 0.25, 0.22
            c = np.array([BOARD_XC, -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)
            # 参考轨迹铺粗：密排蓝球，连起来即一条粗线
            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.008),
                              rgba=Rgba(1.0, 0.36, 0.45, 1.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 True:
        t = root_ctx.get_time()
        if t >= args.duration - 1e-12:
            break
        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 = root_ctx.get_time()

        if t >= next_trail:
            # 末端实际轨迹：定期采样 tool_link 的世界坐标，连成一条黄线画在场景里。
            # 这条线就是「真机跑出来的路径」，可与蓝线（参考轨迹）逐点对照。
            next_trail = t + 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 = t + 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("[faze4_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("[faze4_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("[faze4_real] 常驻中，浏览器打开 %s 查看；Ctrl+C 退出" % meshcat.web_url())
        try:
            while True:
                time.sleep(0.5)
        except KeyboardInterrupt:
            print("[faze4_real] 已退出")


if __name__ == "__main__":
    main()