#!/usr/bin/env python3
"""真机 ROS 录制：订 control_cmd + hand_state，画与仿真相同的三曲线（deg / deg/s / current）。

配合 step 指令（先握中指再张开）：
  python3 tools/ros_real_record_qvt.py --hand left --finger middle --auto-step

手动：终端1 本脚本；终端2 发 ros2 topic pub ...
"""

from __future__ import annotations

import argparse
import json
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
    range_to_arc_left,
    range_to_arc_right,
)

FINGER_CFG = {
    "index": {"ros_idx": 2, "mcp": "index_mcp_pitch", "lo": 0.0, "hi": 1.60},
    "middle": {"ros_idx": 3, "mcp": "middle_mcp_pitch", "lo": 0.0, "hi": 1.60},
    "ring": {"ros_idx": 4, "mcp": "ring_mcp_pitch", "lo": 0.0, "hi": 1.60},
    "pinky": {"ros_idx": 5, "mcp": "pinky_mcp_pitch", "lo": 0.0, "hi": 1.60},
}


def u8_to_rad(u8: float, hand: str, ros_idx: int) -> float:
    pos = [255.0] * 6
    pos[ros_idx] = float(u8)
    if hand == "left":
        arc = range_to_arc_left(pos, "O6")
    else:
        arc = range_to_arc_right(pos, "O6")
    return float(arc[ros_idx])


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("--state-topic", default="")
    parser.add_argument("--info-topic", default="")
    parser.add_argument("--out", type=Path, default=None)
    parser.add_argument("--label", default="ros_real")
    parser.add_argument("--hz", type=float, default=100.0)
    parser.add_argument(
        "--auto-step",
        action="store_true",
        help="自动发 255..255,0,255..255 再张开，录完退出",
    )
    parser.add_argument("--pre-wait", type=float, default=2.0, help="开录后等待(s)")
    parser.add_argument("--hold", type=float, default=3.0, help="每步保持(s)")
    parser.add_argument("--record-after", type=float, default=1.0, help="第二步后多录(s)")
    parser.add_argument(
        "--duration",
        type=float,
        default=0.0,
        help=">0 时录满该秒数自动结束（配合外部 middle_step_cmd_pub）",
    )
    args = parser.parse_args()

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

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

    logs = {
        "t": [],
        "cmd_u8": [],
        "joint_u8": [],
        "cmd_rad": [],
        "q_rad": [],
        "v_rad_s": [],
        "current": [],
    }
    t0 = time.perf_counter()
    running = True
    state_lock = threading.Lock()
    last_cmd_u8 = [255.0] * 6
    last_joint_u8 = [float("nan")] * 6
    last_current = [float("nan")] * 6
    have_state = False
    have_current = False

    class Recorder(Node):
        def __init__(self):
            super().__init__("ros_real_record_qvt")
            self.pub = self.create_publisher(JointState, cmd_topic, 10)
            self.create_subscription(JointState, cmd_topic, self.on_cmd, 50)
            self.create_subscription(JointState, state_topic, self.on_state, 50)
            self.create_subscription(String, info_topic, self.on_info, 10)
            self.create_timer(1.0 / max(1.0, args.hz), self.on_timer)
            self.get_logger().info(
                f"real record | cmd={cmd_topic} | state={state_topic} | "
                f"finger={args.finger} idx={ros_i}"
            )

        def on_cmd(self, msg: JointState):
            nonlocal last_cmd_u8
            pos = list(msg.position)
            if len(pos) >= 6:
                with state_lock:
                    last_cmd_u8 = [float(x) for x in pos[:6]]

        def on_state(self, msg: JointState):
            nonlocal have_state, last_joint_u8, have_current, last_current
            with state_lock:
                pos = list(msg.position)
                if len(pos) >= 6 and all(float(x) >= 0 for x in pos[:6]):
                    have_state = True
                    last_joint_u8 = [float(x) for x in pos[:6]]
                if msg.effort and any(abs(float(x)) > 1e-6 for x in msg.effort):
                    have_current = True
                    last_current = [float(x) for x in msg.effort[:6]]

        def on_info(self, msg: String):
            nonlocal have_current, last_current
            try:
                info = json.loads(msg.data)
                cur = info.get("current")
                if isinstance(cur, list) and len(cur) >= 6:
                    if any(float(x) >= 0 for x in cur[:6]):
                        with state_lock:
                            have_current = True
                            last_current = [float(x) for x in cur[:6]]
            except json.JSONDecodeError:
                pass

        def publish_pose(self, middle_u8: float):
            nonlocal last_cmd_u8
            m = JointState()
            m.position = [255.0, 255.0, 255.0, 255.0, 255.0, 255.0]
            m.position[ros_i] = float(middle_u8)
            with state_lock:
                last_cmd_u8 = list(m.position)
            self.pub.publish(m)

        def hold_pose(self, middle_u8: float, duration: float, rate_hz: float = 10.0):
            """持续重发控制帧（SDK 只响应订阅到的 cmd；单发 --once 可能丢）。"""
            end = time.perf_counter() + duration
            dt = 1.0 / max(1.0, rate_hz)
            while time.perf_counter() < end and running:
                self.publish_pose(middle_u8)
                time.sleep(dt)

        def on_timer(self):
            nonlocal running
            if not running:
                return
            with state_lock:
                joint_u8 = list(last_joint_u8)
                current = list(last_current)
            cmd_u8 = float(last_cmd_u8[ros_i])
            if have_state and not np.isnan(joint_u8[ros_i]):
                q_u8 = float(joint_u8[ros_i])
            else:
                q_u8 = cmd_u8
            t = time.perf_counter() - t0
            q_rad = u8_to_rad(q_u8, args.hand, ros_i)
            cmd_rad = u8_to_rad(cmd_u8, args.hand, ros_i)
            cur = current[ros_i] if len(current) > ros_i else float("nan")
            logs["t"].append(t)
            logs["cmd_u8"].append(cmd_u8)
            logs["joint_u8"].append(q_u8)
            logs["cmd_rad"].append(cmd_rad)
            logs["q_rad"].append(q_rad)
            logs["current"].append(cur)
            if len(logs["t"]) >= 2:
                dt = logs["t"][-1] - logs["t"][-2]
                if dt > 1e-9:
                    v = (logs["q_rad"][-1] - logs["q_rad"][-2]) / dt
                else:
                    v = 0.0
            else:
                v = 0.0
            logs["v_rad_s"].append(v)

    rclpy.init()
    node = Recorder()

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

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

    if args.auto_step:
        deadline = time.perf_counter() + max(5.0, args.pre_wait)
        while time.perf_counter() < deadline and not have_state:
            time.sleep(0.05)
        if not have_state:
            node.get_logger().warn(
                "仍未收到 hand_state（可能 SDK 未跑 / Domain 不对 / 全是 -1）"
            )
        time.sleep(max(0.5, args.pre_wait))
        node.get_logger().info("step1: middle close (u8=0)")
        node.hold_pose(0.0, args.hold)
        node.get_logger().info("step2: middle open (u8=255)")
        node.hold_pose(255.0, args.hold + args.record_after)
        running = False
    else:
        node.get_logger().info(
            "listen mode: 外部发令 (middle_step_cmd_pub) 或 Ctrl+C 结束"
        )
        t_end = (
            time.perf_counter() + args.duration if args.duration > 0 else None
        )
        try:
            while running:
                if t_end is not None and time.perf_counter() >= t_end:
                    break
                time.sleep(0.05)
        except KeyboardInterrupt:
            running = False

    time.sleep(0.2)
    node.destroy_node()
    if rclpy.ok():
        rclpy.shutdown()

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

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

    if len(logs["v_rad_s"]) >= 5:
        logs["v_rad_s"] = np.convolve(
            logs["v_rad_s"], np.ones(5) / 5.0, mode="same"
        )

    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["joint_u8"],
                logs["cmd_rad"],
                logs["q_rad"],
                logs["v_rad_s"],
                logs["current"],
            ]
        ),
        delimiter=",",
        header="t_s,cmd_u8,joint_u8,cmd_rad,q_rad,v_rad_s,current",
        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

    mcp_lo, mcp_hi = cfg["lo"], cfg["hi"]
    fig, axes = plt.subplots(3, 1, figsize=(10, 8), sharex=True)
    fig.suptitle(
        f"ROS->Real {args.hand} {cfg['mcp']}  "
        f"range=[{mcp_lo:.2f}, {mcp_hi:.2f}] rad  "
        f"(255=open/0deg, 0=close/~92deg)",
        fontsize=11,
    )
    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_rad"]), "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_rad_s"]), "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)

    if have_current and np.any(np.isfinite(logs["current"])):
        axes[2].plot(
            logs["t"], logs["current"], "C3", lw=1.6, label="current (hand_info)"
        )
        axes[2].set_ylabel("current")
    else:
        axes[2].text(
            0.5,
            0.5,
            "no current (subscribe hand_info or state.effort)",
            ha="center",
            va="center",
            transform=axes[2].transAxes,
        )
        axes[2].set_ylabel("current")
    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'])}  state={'yes' if have_state else 'NO'}")
    cmd_span = float(np.max(logs["cmd_u8"]) - np.min(logs["cmd_u8"]))
    joint_span = float(np.max(logs["joint_u8"]) - np.min(logs["joint_u8"]))
    print(f"cmd_u8 span={cmd_span:.1f}  joint_u8 span={joint_span:.1f}")
    if cmd_span < 1.0:
        print("[warn] cmd 全程几乎不变 → 控制指令可能没发出去或未写入录制")
    if not have_state:
        print("[warn] 未收到有效 state → q 只能用 cmd 代替，且真机可能没动")
    if cmd_span >= 1.0 and joint_span < 3.0 and have_state:
        print("[warn] 有 cmd 变化但 joint 几乎不动 → 查 CAN / SDK / 力矩限制")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())