#!/usr/bin/env python3
"""真机 vs 仿真 q/v/力矩(或电流) 对比图。

时间原点 = 首次收到「握紧」控制指令 (cmd_u8: 255→0) 的时刻 − 1 s。

用法:
  python3 tools/plot_real_sim_compare.py \\
    reports/O6_middle_ros/ros_real_middle_left_20260716_172221.csv \\
    reports/O6_middle_ros/ros_sim_middle_left_20260716_114515.csv
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np


def load_csv(path: Path):
    hdr = path.read_text().splitlines()[0].strip().split(",")
    data = np.loadtxt(path, delimiter=",", skiprows=1)
    col = {name: i for i, name in enumerate(hdr)}
    return hdr, col, data


def first_close_cmd_time(t: np.ndarray, cmd_u8: np.ndarray) -> float:
    """首次 cmd 从 255 阶跃到 0（握紧）的时刻。"""
    diff = np.diff(cmd_u8)
    idx = np.where(diff < -100)[0]
    if len(idx) == 0:
        raise ValueError("未检测到 cmd 255→0 阶跃")
    return float(t[idx[0] + 1])


def align_series(t, *arrays, t_cmd: float, t_pre: float = 1.0):
    t0 = t_cmd - t_pre
    t_rel = t - t0
    return (t_rel, *arrays)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("real_csv", type=Path)
    parser.add_argument("sim_csv", type=Path)
    parser.add_argument("-o", "--out", type=Path, default=None)
    parser.add_argument("--t-pre", type=float, default=1.0)
    parser.add_argument("--t-max", type=float, default=8.0, help="相对原点后的最大时间(s)")
    args = parser.parse_args()

    _, cr, dr = load_csv(args.real_csv)
    _, cs, ds = load_csv(args.sim_csv)

    t_r, cmd_r = dr[:, cr["t_s"]], dr[:, cr["cmd_u8"]]
    t_s, cmd_s = ds[:, cs["t_s"]], ds[:, cs["cmd_u8"]]

    t_cmd_r = first_close_cmd_time(t_r, cmd_r)
    t_cmd_s = first_close_cmd_time(t_s, cmd_s)

    q_r = dr[:, cr["q_rad"]]
    cmd_rad_r = dr[:, cr["cmd_rad"]]
    v_r = dr[:, cr["v_rad_s"]]
    q_s = ds[:, cs["q_rad"]]
    cmd_rad_s = ds[:, cs["cmd_rad"]]
    v_s = ds[:, cs["v_rad_s"]]

    cur_r = dr[:, cr["current"]] if "current" in cr else np.full_like(q_r, np.nan)
    tau_s = ds[:, cs["tau_Nm"]] if "tau_Nm" in cs else np.full_like(q_s, np.nan)

    tr, qr, cmdr, vr, cur_r = align_series(
        t_r, q_r, cmd_rad_r, v_r, cur_r, t_cmd=t_cmd_r, t_pre=args.t_pre
    )
    ts, qs, cmds, vs, tau_s = align_series(
        t_s, q_s, cmd_rad_s, v_s, tau_s, t_cmd=t_cmd_s, t_pre=args.t_pre
    )

    mask_r = (tr >= -0.05) & (tr <= args.t_max)
    mask_s = (ts >= -0.05) & (ts <= args.t_max)

    out = args.out
    if out is None:
        out = args.real_csv.parent / "ros_real_vs_sim_middle_left_compare_qvt.png"

    import os

    root = Path(__file__).resolve().parents[1]
    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=(11, 9), sharex=True)
    fig.suptitle(
        "O6 middle_mcp_pitch  real vs sim\n"
        f"t=0 = {args.t_pre:.0f}s before first close command (cmd_u8: 255->0)",
        fontsize=12,
    )

    # --- angle ---
    ax = axes[0]
    # 控制指令：真机/仿真时序一致，只画一条黑色虚线
    cmd_t = ts[mask_s] if np.any(mask_s) else tr[mask_r]
    cmd_y = np.rad2deg(cmds[mask_s]) if np.any(mask_s) else np.rad2deg(cmdr[mask_r])
    ax.plot(cmd_t, cmd_y, color="black", ls="--", lw=1.2, label="cmd")
    ax.plot(ts[mask_s], np.rad2deg(qs[mask_s]), color="#1f77b4", lw=1.8, label="sim q")
    ax.plot(tr[mask_r], np.rad2deg(qr[mask_r]), color="#ff7f0e", lw=1.8, label="real q")
    ax.axvline(1.0, color="gray", ls=":", lw=1.0, alpha=0.8)
    ax.set_ylabel("angle (deg)")
    ax.legend(loc="best", fontsize=9)
    ax.grid(True, alpha=0.3)

    # --- velocity ---
    ax = axes[1]
    ax.plot(ts[mask_s], np.rad2deg(vs[mask_s]), color="#1f77b4", lw=1.8, label="sim v")
    ax.plot(tr[mask_r], np.rad2deg(vr[mask_r]), color="#ff7f0e", lw=1.8, label="real v")
    ax.axvline(1.0, color="gray", ls=":", lw=1.0, alpha=0.8)
    ax.set_ylabel("velocity (deg/s)")
    ax.legend(loc="best")
    ax.grid(True, alpha=0.3)

    # --- effort: sim tau vs real current ---
    ax = axes[2]
    ax.plot(ts[mask_s], tau_s[mask_s], color="#1f77b4", lw=1.8, label="sim τ (N·m)")
    if np.any(np.isfinite(cur_r[mask_r])):
        ax.plot(tr[mask_r], cur_r[mask_r], color="#ff7f0e", lw=1.8, label="real current")
        ax.set_ylabel("τ (N·m) / current")
    else:
        ax.plot([], [], color="#ff7f0e", label="real current (no data)")
        ax.set_ylabel("torque (N·m)")
        ax.text(
            0.5,
            0.5,
            "real current not recorded",
            transform=ax.transAxes,
            ha="center",
            va="center",
            fontsize=10,
            color="#666",
        )
    ax.axvline(1.0, color="gray", ls=":", lw=1.0, alpha=0.8)
    ax.set_xlabel("time (s)  [t=0: 1s before first close cmd]")
    ax.legend(loc="best")
    ax.grid(True, alpha=0.3)

    fig.tight_layout()
    fig.savefig(out, dpi=150)
    plt.close(fig)

    print(f"plot -> {out}")
    print(f"real cmd↓ at abs t={t_cmd_r:.3f}s  -> origin t0={t_cmd_r - args.t_pre:.3f}s")
    print(f"sim  cmd↓ at abs t={t_cmd_s:.3f}s  -> origin t0={t_cmd_s - args.t_pre:.3f}s")


if __name__ == "__main__":
    main()