#!/usr/bin/env python3
"""O6 单指位置控制：从最小扫到最大，记录并绘制 角度/速度/力矩。"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import mujoco
import mujoco.viewer
import numpy as np

ROOT = Path(__file__).resolve().parents[2]
DEFAULT_XML = (
    ROOT
    / "src/linkerhand-sim/linker_hand_mujoco_ros2/linker_hand_mujoco_ros2"
    / "urdf/O6/linker_hand_o6_left/linker_hand_o6_left.xml"
)

# 四指：根节 mcp_pitch + 远端 dip（mimic 0.89）
# ROS 6 维索引：0拇指弯 1拇指横摆 2食指 3中指 4无名指 5小指
FINGER_CFG = {
    "index": {
        "mcp": "index_mcp_pitch",
        "dip": "index_dip",
        "mcp_act": "index_mcp_pitch_pos",
        "dip_act": "index_dip_pos",
        "mimic": 0.89,
        "ros_idx": 2,
    },
    "middle": {
        "mcp": "middle_mcp_pitch",
        "dip": "middle_dip",
        "mcp_act": "middle_mcp_pitch_pos",
        "dip_act": "middle_dip_pos",
        "mimic": 0.89,
        "ros_idx": 3,
    },
    "ring": {
        "mcp": "ring_mcp_pitch",
        "dip": "ring_dip",
        "mcp_act": "ring_mcp_pitch_pos",
        "dip_act": "ring_dip_pos",
        "mimic": 0.89,
        "ros_idx": 4,
    },
    "pinky": {
        "mcp": "pinky_mcp_pitch",
        "dip": "pinky_dip",
        "mcp_act": "pinky_mcp_pitch_pos",
        "dip_act": "pinky_dip_pos",
        "mimic": 0.89,
        "ros_idx": 5,
    },
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--xml", type=Path, default=DEFAULT_XML)
    parser.add_argument("--hand", choices=["left", "right"], default="left")
    parser.add_argument(
        "--finger",
        choices=list(FINGER_CFG.keys()),
        default="middle",
        help="要扫的手指（默认 middle）",
    )
    parser.add_argument("--duration", type=float, default=8.0, help="单向扫过时长 (s)")
    parser.add_argument("--out", type=Path, default=None)
    parser.add_argument("--viewer", action="store_true", help="打开 MuJoCo viewer")
    args = parser.parse_args()

    cfg = FINGER_CFG[args.finger]
    if args.out is None:
        args.out = ROOT / f"reports/O6_{args.finger}_sweep"

    if args.hand == "right":
        args.xml = (
            ROOT
            / "src/linkerhand-sim/linker_hand_mujoco_ros2/linker_hand_mujoco_ros2"
            / "urdf/O6/linker_hand_o6_right/linker_hand_o6_right.xml"
        )

    args.out.mkdir(parents=True, exist_ok=True)
    model = mujoco.MjModel.from_xml_path(str(args.xml))
    data = mujoco.MjData(model)

    mcp_jnt = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, cfg["mcp"])
    dip_jnt = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, cfg["dip"])
    mcp_act = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, cfg["mcp_act"])
    dip_act = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, cfg["dip_act"])
    if min(mcp_jnt, dip_jnt, mcp_act, dip_act) < 0:
        raise RuntimeError(f"找不到 {args.finger} joint/actuator，请检查 MJCF")

    mcp_qadr = int(model.jnt_qposadr[mcp_jnt])
    mcp_dadr = int(model.jnt_dofadr[mcp_jnt])
    mcp_lo, mcp_hi = map(float, model.jnt_range[mcp_jnt])
    dip_lo, dip_hi = map(float, model.jnt_range[dip_jnt])
    mimic = float(cfg["mimic"])
    dt = float(model.opt.timestep)

    half = args.duration
    total_t = 2.0 * half
    n_steps = int(total_t / dt)
    t_arr = np.arange(n_steps) * dt

    def cmd_at(t: float) -> float:
        if t <= half:
            return mcp_lo + (mcp_hi - mcp_lo) * (t / half)
        u = (t - half) / half
        return mcp_hi + (mcp_lo - mcp_hi) * u

    data.ctrl[:] = 0.0
    mujoco.mj_forward(model, data)

    logs = {
        "t": [],
        "cmd_mcp": [],
        "q_mcp": [],
        "v_mcp": [],
        "tau_mcp": [],
        "q_dip": [],
        "tau_dip": [],
    }

    viewer = None
    if args.viewer:
        viewer = mujoco.viewer.launch_passive(model, data)

    for i in range(n_steps):
        t = t_arr[i]
        q_cmd = cmd_at(t)
        dip_cmd = float(np.clip(q_cmd * mimic, dip_lo, dip_hi))
        data.ctrl[mcp_act] = q_cmd
        data.ctrl[dip_act] = dip_cmd
        mujoco.mj_step(model, data)

        logs["t"].append(t)
        logs["cmd_mcp"].append(q_cmd)
        logs["q_mcp"].append(float(data.qpos[mcp_qadr]))
        logs["v_mcp"].append(float(data.qvel[mcp_dadr]))
        logs["tau_mcp"].append(float(data.actuator_force[mcp_act]))
        logs["q_dip"].append(float(data.qpos[int(model.jnt_qposadr[dip_jnt])]))
        logs["tau_dip"].append(float(data.actuator_force[dip_act]))

        if viewer is not None:
            viewer.sync()

    if viewer is not None:
        viewer.close()

    for k in logs:
        logs[k] = np.asarray(logs[k], dtype=float)

    csv_path = args.out / f"{args.finger}_sweep_{args.hand}.csv"
    header = "t_s,cmd_mcp_rad,q_mcp_rad,v_mcp_rad_s,tau_mcp_Nm,q_dip_rad,tau_dip_Nm"
    np.savetxt(
        csv_path,
        np.column_stack(
            [
                logs["t"],
                logs["cmd_mcp"],
                logs["q_mcp"],
                logs["v_mcp"],
                logs["tau_mcp"],
                logs["q_dip"],
                logs["tau_dip"],
            ]
        ),
        delimiter=",",
        header=header,
        comments="",
    )

    fig, axes = plt.subplots(3, 1, figsize=(10, 8), sharex=True)
    fig.suptitle(
        f"O6 {args.hand} {cfg['mcp']} position sweep  "
        f"range=[{mcp_lo:.2f}, {mcp_hi:.2f}] rad",
        fontsize=13,
    )

    axes[0].plot(logs["t"], np.rad2deg(logs["cmd_mcp"]), "k--", lw=1.2, label="cmd")
    axes[0].plot(logs["t"], np.rad2deg(logs["q_mcp"]), "C0", lw=1.8, label="q")
    axes[0].set_ylabel("angle (deg)")
    axes[0].legend(loc="best")
    axes[0].grid(True, alpha=0.3)

    axes[1].plot(logs["t"], np.rad2deg(logs["v_mcp"]), "C1", lw=1.8, label="v")
    axes[1].set_ylabel("velocity (deg/s)")
    axes[1].legend(loc="best")
    axes[1].grid(True, alpha=0.3)

    axes[2].plot(logs["t"], logs["tau_mcp"], "C3", lw=1.8, label="tau (actuator_force)")
    axes[2].set_ylabel("torque (N·m)")
    axes[2].set_xlabel("time (s)")
    axes[2].legend(loc="best")
    axes[2].grid(True, alpha=0.3)

    fig.tight_layout()
    png_path = args.out / f"{args.finger}_sweep_{args.hand}_qvt.png"
    fig.savefig(png_path, dpi=140)
    plt.close(fig)

    # ROS 示例：只动该指
    pos = [255.0] * 6
    pos_open = pos.copy()
    pos_close = pos.copy()
    pos_close[cfg["ros_idx"]] = 0.0

    print("=" * 60)
    print(f"finger: {args.finger}  joint: {cfg['mcp']}")
    print(f"model:  {args.xml}")
    print(
        f"range:  [{mcp_lo:.3f}, {mcp_hi:.3f}] rad "
        f"= [{np.rad2deg(mcp_lo):.1f}, {np.rad2deg(mcp_hi):.1f}] deg"
    )
    print(f"mimic:  dip = mcp * {mimic}")
    print(f"sweep:  {mcp_lo:.3f} -> {mcp_hi:.3f} -> {mcp_lo:.3f}  total {total_t:.1f}s")
    print(f"CSV:    {csv_path}")
    print(f"plot:   {png_path}")
    print(
        "peak |q|={:.3f} rad  |v|={:.3f} rad/s  |tau|={:.3f} N·m".format(
            np.max(np.abs(logs["q_mcp"])),
            np.max(np.abs(logs["v_mcp"])),
            np.max(np.abs(logs["tau_mcp"])),
        )
    )
    print("=" * 60)
    print(f"ROS 只弯{args.finger}（索引 {cfg['ros_idx']}）握紧示例 position:")
    print(f"  {pos_close}")
    print(f"张开示例: {pos_open}")


if __name__ == "__main__":
    main()