#!/usr/bin/env python3
"""ROS 控制指令 → 本地 MuJoCo → 记录并绘制 角度/速度/力矩。

订 /cb_<hand>_hand_control_cmd（0～255），用与仿真节点相同的 O6 映射写 ctrl，
从 MuJoCo 读 q / v / actuator_force，Ctrl+C 存 CSV + 三曲线图。

用法：
  # 终端1（可先不开 run_sim.sh，本脚本自带 MuJoCo）
  python3 tools/ros_mujoco_record_qvt.py --hand left --finger middle --viewer

  # 终端2 发 ROS 指令，或加 --self-sweep 自动发中指慢扫
  python3 tools/ros_mujoco_record_qvt.py --hand left --finger middle --viewer --self-sweep
"""

from __future__ import annotations

import argparse
import sys
import threading
import time
from datetime import datetime
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
PKG = (
    ROOT
    / "src/linkerhand-sim/linker_hand_mujoco_ros2/linker_hand_mujoco_ros2"
)
sys.path.insert(0, str(PKG.parent))

from linker_hand_mujoco_ros2.utils.mapping import (  # noqa: E402
    O6_JOINT_MAP,
    O6_MIMIC,
    apply_mimic,
    range_to_arc_left,
    range_to_arc_right,
)

FINGER_CFG = {
    "index": {
        "mcp": "index_mcp_pitch",
        "mcp_act": "index_mcp_pitch_pos",
        "ros_idx": 2,
    },
    "middle": {
        "mcp": "middle_mcp_pitch",
        "mcp_act": "middle_mcp_pitch_pos",
        "ros_idx": 3,
    },
    "ring": {
        "mcp": "ring_mcp_pitch",
        "mcp_act": "ring_mcp_pitch_pos",
        "ros_idx": 4,
    },
    "pinky": {
        "mcp": "pinky_mcp_pitch",
        "mcp_act": "pinky_mcp_pitch_pos",
        "ros_idx": 5,
    },
}


def map_position_array(position, joint_map):
    mapped = [0.0] * len(joint_map)
    for target_idx, source_idx in joint_map.items():
        if source_idx < len(position):
            mapped[target_idx] = position[source_idx]
    return mapped


def u8_to_ctrl(position_u8, hand_type: str):
    if hand_type == "left":
        tmp = range_to_arc_left(position_u8, "O6")
    else:
        tmp = range_to_arc_right(position_u8, "O6")
    res = map_position_array(tmp, O6_JOINT_MAP)
    from linker_hand_mujoco_ros2.utils.mapping import O6_JOINT_ARC

    res = apply_mimic(res, O6_MIMIC, O6_JOINT_ARC)
    return res


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--hand", choices=["left", "right"], default="left")
    parser.add_argument("--finger", choices=list(FINGER_CFG.keys()), default="middle")
    parser.add_argument("--cmd-topic", default="")
    parser.add_argument("--out", type=Path, default=None)
    parser.add_argument("--viewer", action="store_true")
    parser.add_argument(
        "--self-sweep",
        action="store_true",
        help="本进程自动往控制话题发 255→0→255 慢扫（方便自测）",
    )
    parser.add_argument(
        "--exit-after-sweep",
        action="store_true",
        help="配合 --self-sweep：扫完自动结束并画图",
    )
    parser.add_argument("--duration", type=float, default=8.0, help="self-sweep 单向时长")
    parser.add_argument(
        "--record-duration",
        type=float,
        default=0.0,
        help=">0 时录满该秒数自动结束（配合 middle_step_cmd_pub）",
    )
    parser.add_argument("--label", default="ros_sim")
    args = parser.parse_args()
    if args.self_sweep and not args.exit_after_sweep:
        # 自测默认扫完就退出，避免一直挂着
        args.exit_after_sweep = True

    cfg = FINGER_CFG[args.finger]
    cmd_topic = args.cmd_topic or f"/cb_{args.hand}_hand_control_cmd"
    out_dir = args.out or (ROOT / f"reports/O6_{args.finger}_ros")
    out_dir.mkdir(parents=True, exist_ok=True)

    xml = (
        PKG
        / f"urdf/O6/linker_hand_o6_{args.hand}/linker_hand_o6_{args.hand}.xml"
    )

    import mujoco
    import mujoco.viewer

    model = mujoco.MjModel.from_xml_path(str(xml))
    data = mujoco.MjData(model)
    model.opt.disableflags = 0

    mcp_jnt = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, cfg["mcp"])
    mcp_act = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, cfg["mcp_act"])
    if mcp_jnt < 0 or mcp_act < 0:
        raise RuntimeError(f"joint/actuator not found for {args.finger}")
    qadr = int(model.jnt_qposadr[mcp_jnt])
    dadr = int(model.jnt_dofadr[mcp_jnt])
    mcp_lo, mcp_hi = map(float, model.jnt_range[mcp_jnt])

    ctrl_lock = threading.Lock()
    ctrl_values = np.zeros(model.nu)
    last_cmd_u8 = [255.0] * 6
    running = True
    sweep_done = threading.Event()
    logs = {"t": [], "cmd_rad": [], "q": [], "v": [], "tau": [], "cmd_u8": []}
    t0 = time.perf_counter()
    t_stop = (
        t0 + args.record_duration if args.record_duration > 0 else None
    )

    import signal

    def _stop(*_):
        nonlocal running
        running = False

    signal.signal(signal.SIGINT, _stop)
    signal.signal(signal.SIGTERM, _stop)

    import rclpy
    from rclpy.node import Node
    from sensor_msgs.msg import JointState

    class CmdNode(Node):
        def __init__(self):
            super().__init__("ros_mujoco_record_qvt")
            self.create_subscription(JointState, cmd_topic, self.on_cmd, 50)
            self.get_logger().info(
                f"MuJoCo+ROS record | topic={cmd_topic} | finger={args.finger} | "
                f"joint={cfg['mcp']} range=[{mcp_lo:.2f},{mcp_hi:.2f}]"
            )
            self.get_logger().info("Ctrl+C to stop and plot")

        def on_cmd(self, msg: JointState):
            nonlocal last_cmd_u8
            pos = list(msg.position)
            if len(pos) < 6:
                pos = pos + [255.0] * (6 - len(pos))
            last_cmd_u8 = pos[:6]
            mapped = u8_to_ctrl(last_cmd_u8, args.hand)
            with ctrl_lock:
                n = min(len(mapped), model.nu)
                ctrl_values[:n] = mapped[:n]

    rclpy.init()
    node = CmdNode()

    def spin_ros():
        while running and rclpy.ok():
            rclpy.spin_once(node, timeout_sec=0.01)

    ros_thread = threading.Thread(target=spin_ros, daemon=True)
    ros_thread.start()

    pub = None
    if args.self_sweep:
        pub = node.create_publisher(JointState, cmd_topic, 10)
        time.sleep(0.3)

        def sweep_pub():
            nonlocal running
            half = args.duration
            ros_i = cfg["ros_idx"]
            # 255 -> 0
            t_start = time.time()
            while running and time.time() - t_start < half:
                u = (time.time() - t_start) / half
                mid = 255.0 * (1.0 - u)
                m = JointState()
                m.position = [255.0, 255.0, 255.0, 255.0, 255.0, 255.0]
                m.position[ros_i] = mid
                pub.publish(m)
                time.sleep(0.01)
            # 0 -> 255
            t_start = time.time()
            while running and time.time() - t_start < half:
                u = (time.time() - t_start) / half
                mid = 255.0 * u
                m = JointState()
                m.position = [255.0, 255.0, 255.0, 255.0, 255.0, 255.0]
                m.position[ros_i] = mid
                pub.publish(m)
                time.sleep(0.01)
            print("self-sweep finished", flush=True)
            sweep_done.set()
            if args.exit_after_sweep:
                running = False

        threading.Thread(target=sweep_pub, daemon=True).start()

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

    try:
        while running:
            if viewer is not None and not viewer.is_running():
                break
            if t_stop is not None and time.perf_counter() >= t_stop:
                break
            with ctrl_lock:
                data.ctrl[:] = ctrl_values
            mujoco.mj_step(model, data)
            if viewer is not None:
                viewer.sync()

            t = time.perf_counter() - t0
            q = float(data.qpos[qadr])
            v = float(data.qvel[dadr])
            tau = float(data.actuator_force[mcp_act])
            cmd_rad = float(ctrl_values[mcp_act]) if mcp_act < len(ctrl_values) else 0.0
            logs["t"].append(t)
            logs["cmd_rad"].append(cmd_rad)
            logs["q"].append(q)
            logs["v"].append(v)
            logs["tau"].append(tau)
            logs["cmd_u8"].append(float(last_cmd_u8[cfg["ros_idx"]]))
            time.sleep(0.001)
    except KeyboardInterrupt:
        pass
    finally:
        running = False
        if viewer is not None:
            viewer.close()
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()

    if len(logs["t"]) < 2:
        print("no samples")
        return

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

    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    stem = f"{args.label}_{args.finger}_{args.hand}_{stamp}"
    csv_path = out_dir / f"{stem}.csv"
    png_path = out_dir / f"{stem}_qvt.png"

    np.savetxt(
        csv_path,
        np.column_stack(
            [
                logs["t"],
                logs["cmd_u8"],
                logs["cmd_rad"],
                logs["q"],
                logs["v"],
                logs["tau"],
            ]
        ),
        delimiter=",",
        header="t_s,cmd_u8,cmd_rad,q_rad,v_rad_s,tau_Nm",
        comments="",
    )

    import os

    os.environ.setdefault("MPLCONFIGDIR", str(ROOT / ".mplconfig"))
    Path(os.environ["MPLCONFIGDIR"]).mkdir(parents=True, exist_ok=True)
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(3, 1, figsize=(10, 8), sharex=True)
    fig.suptitle(
        f"ROS->MuJoCo {args.hand} {cfg['mcp']}  "
        f"range=[{mcp_lo:.2f}, {mcp_hi:.2f}] rad",
        fontsize=12,
    )
    axes[0].plot(logs["t"], np.rad2deg(logs["cmd_rad"]), "k--", lw=1.2, label="cmd")
    axes[0].plot(logs["t"], np.rad2deg(logs["q"]), "C0", lw=1.6, 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"]), "C1", lw=1.6, 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"], "C3", lw=1.6, 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()
    fig.savefig(png_path, dpi=140)
    plt.close(fig)

    print(f"CSV  -> {csv_path}")
    print(f"plot -> {png_path}")
    print(f"samples={len(logs['t'])}")


if __name__ == "__main__":
    main()