# -*- coding: utf-8 -*-
"""
测试脚本：验证 InteractionEstimator + HapticRenderer 的正确性，
并绘制论文中使用的误差和能量罐曲线。

运行：
    python test_haptic_render.py

依赖：
    - pinocchio
    - numpy
    - matplotlib
    - omegaconf
"""

import numpy as np
import pinocchio as pin
from omegaconf import OmegaConf
import matplotlib.pyplot as plt

import os, sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from core.interaction_estimater import InteractionEstimator
from core.haptic_render import HapticRenderer, TankParams


def moving_average(x, window: int = 7):
    """简单滑动平均，用于平滑 alpha 等曲线。"""
    if window <= 1:
        return x
    kernel = np.ones(window, dtype=float) / float(window)
    # 使用 same 保持长度一致
    return np.convolve(x, kernel, mode="same")


def build_model_with_frames():
    """
    备用：构造一个 6-DoF 示例机械臂，并添加：
      - chest_link: 挂在关节1
      - ee_link:    挂在末端关节
    目前脚本直接从 URDF 加载，不一定用得到。
    """
    model = pin.buildSampleModelManipulator()
    chest_joint_id = 1
    chest_frame_name = "chest_link"
    model.addFrame(pin.Frame(
        chest_frame_name,
        chest_joint_id,
        chest_joint_id,
        pin.SE3.Identity(),
        pin.FrameType.OP_FRAME
    ))
    ee_joint_id = model.njoints - 1
    ee_frame_name = "ee_link"
    model.addFrame(pin.Frame(
        ee_frame_name,
        ee_joint_id,
        ee_joint_id,
        pin.SE3.Identity(),
        pin.FrameType.OP_FRAME
    ))
    return model, chest_frame_name, ee_frame_name


def main():
    conf = OmegaConf.load("./config/config.yaml")
    np.random.seed(0)

    # 1) 构造主/从模型（这里用真实 URDF）
    master_model, _, _ = pin.buildModelsFromUrdf(str(conf.master_urdf))
    slave_model,  _, _ = pin.buildModelsFromUrdf(str(conf.slave_urdf))

    # 从端交互估计器
    est_slave = InteractionEstimator(
        slave_model,
        chest_frame_name=conf.s_base_frame,
        ee_frame_name=conf.s_ee_frame,
        lambda_damp=conf.interaction_est.lambda_damp,
    )

    # 主端渲染器
    renderer = HapticRenderer(
        master_model,
        chest_frame_name=conf.m_base_frame,
        ee_frame_name=conf.m_ee_frame,
        feedback_strength=conf.haptic_render.feedback_strength,
        E_init=conf.haptic_render.E_init,
        E_max=conf.haptic_render.E_max,
        alpha_floor=conf.haptic_render.alpha_floor,
        alpha_ceil=conf.haptic_render.alpha_ceil,
        E0=conf.haptic_render.E0,
    )

    nq_s, nv_s = slave_model.nq, slave_model.nv
    nq_m, nv_m = master_model.nq, master_model.nv

    dt = 0.002
    n_steps = 500  # 足够画出平滑曲线

    print("========== TEST START ==========")
    print(f"slave nq={nq_s}, master nq={nq_m}, dt={dt}s, steps={n_steps}")
    print("--------------------------------")

    # -------- 统计量 / 曲线数据 --------
    times = []
    err_tau_int_hist = []
    err_CF_hist = []
    err_tau_fb_hist = []
    E_hist = []
    alpha_hist = []

    max_err_tau_int = 0.0
    max_err_CF      = 0.0
    max_err_tau_fb  = 0.0

    for k in range(n_steps):
        t = k * dt
        times.append(t)

        # 2) 随机生成从端状态（可改成真实记录或特定轨迹）
        q_s  = 0.2 * (np.random.rand(nq_s) - 0.5)   # [-0.1,0.1]
        qd_s = 0.1 * (np.random.rand(nv_s) - 0.5)
        qdd_s = np.zeros_like(qd_s)

        # 从端模型扭矩（M qdd + C qd + g）
        tau_model_s = est_slave._tau_model(q_s, qd_s, qdd_s, tau_ff_fric=None)

        # 从端胸腔雅可比 C J_s
        CJ_s = est_slave._chest_jacobian(q_s, qd_s)  # 6 x nv_s

        # 3) 人为构造“真实”交互扳手 C F_true
        #    也可以改成随时间变化的模式，例如正弦。
        CF_true = np.array([10.0, 0.0, 0.0,   # 10 N 沿 Cx
                            0.0,  0.0, 0.0])

        # 真 · 交互关节力矩
        tau_int_true = CJ_s.T @ CF_true

        # 构造测量力矩：tau_meas = tau_model + tau_int_true
        tau_meas_s = tau_model_s + tau_int_true

        # 4) 调用 InteractionEstimator 估计
        tau_int_est, CF_int_est, CJ_s_est = est_slave.estimate(
            q_s, qd_s, qdd_s,
            tau_meas=tau_meas_s,
            tau_ff_fric=None
        )

        # 从端末端在 C 系的速度扭量（用于能量罐功率）
        V_slave_C = CJ_s @ qd_s

        # 5) 主端状态（此处用从端状态代替，仅做算法验证）
        q_m  = q_s.copy()
        qd_m = qd_s.copy()

        tau_fb_m_est, alpha = renderer.render_from_CF(
            q_m, qd_m,
            CF_int_slave_C=CF_int_est,
            V_slave_C=V_slave_C,
            dt=dt
        )

        # 主端胸腔雅可比，用于构造“真 · 反馈扭矩”（不经过能量罐）
        CJ_m = renderer.CJ_master.chest_jacobian(q_m, qd_m)
        tau_fb_m_true = CJ_m.T @ CF_true

        # -------- 误差计算 --------
        err_tau_int = np.linalg.norm(tau_int_est - tau_int_true)
        err_CF      = np.linalg.norm(CF_int_est - CF_true)
        err_tau_fb  = np.linalg.norm(tau_fb_m_est - tau_fb_m_true)

        err_tau_int_hist.append(err_tau_int)
        err_CF_hist.append(err_CF)
        err_tau_fb_hist.append(err_tau_fb)

        max_err_tau_int = max(max_err_tau_int, err_tau_int)
        max_err_CF      = max(max_err_CF,      err_CF)
        max_err_tau_fb  = max(max_err_tau_fb,  err_tau_fb)

        # 能量罐状态
        E_hist.append(renderer.tank.E)
        alpha_hist.append(alpha)

        # ---- 每隔若干步打印一次 ----
        if k % 100 == 0:
            print(f"\n--- Step {k} (t = {t:.3f} s) ---")
            print(f"‣ ||tau_int_true|| = {np.linalg.norm(tau_int_true):.4e}")
            print(f"  ||tau_int_est - tau_int_true|| = {err_tau_int:.4e}")
            print(f"‣ ||CF_true|| = {np.linalg.norm(CF_true):.4e}")
            print(f"  ||CF_int_est - CF_true||       = {err_CF:.4e}")
            print(f"‣ ||tau_fb_m_true|| = {np.linalg.norm(tau_fb_m_true):.4e}")
            print(f"  ||tau_fb_m_est - tau_fb_m_true|| = {err_tau_fb:.4e}")
            print(f"  Energy tank: E = {renderer.tank.E:.4f}, alpha = {alpha:.4f}")

    # ---------- 数值结果摘要 ----------
    print("\n========== SUMMARY ==========")
    print(f"max ||tau_int_est - tau_int_true|| = {max_err_tau_int:.4e}")
    print(f"max ||CF_int_est - CF_true||       = {max_err_CF:.4e}")
    print(f"max ||tau_fb_est - tau_fb_true||   = {max_err_tau_fb:.4e}")

    tol_tau_int = 1e-3
    tol_CF      = 1e-3
    tol_tau_fb  = 1e-3
    if max_err_tau_int < tol_tau_int and max_err_CF < tol_CF and max_err_tau_fb < tol_tau_fb:
        print("[PASS] 所有误差均在容许范围内。")
    else:
        print("[WARN] 误差超出阈值，请检查算法或考虑调小阻尼 lambda_damp。")

    # =====================================================
    #              绘图部分（适合论文呈现）
    # =====================================================

    times = np.array(times)
    err_tau_int_hist = np.array(err_tau_int_hist)
    err_CF_hist      = np.array(err_CF_hist)
    err_tau_fb_hist  = np.array(err_tau_fb_hist)
    E_hist           = np.array(E_hist)
    alpha_hist       = np.array(alpha_hist)

    # 对 alpha 做轻微平滑（论文图更清晰）
    alpha_smooth = moving_average(alpha_hist, window=7)

    # ------------------ 误差子图（3 个独立子图） ------------------
    fig_err, (ax1, ax2, ax3) = plt.subplots(
        3, 1, sharex=True, figsize=(6, 7)
    )

    # (a) Interaction torque estimation error
    ax1.plot(
        times,
        err_tau_int_hist,
        color="C0",
        label=r"$\|\tau_{\mathrm{int}}^{\mathrm{est}}-\tau_{\mathrm{int}}^{\mathrm{true}}\|$",
    )
    ax1.set_title("(a) Interaction torque estimation error")
    ax1.set_ylabel(r"Error norm $\|\cdot\|$")
    ax1.set_yscale("log")
    ax1.grid(True, linestyle="--", linewidth=0.5)
    ax1.legend(loc="lower right")

    # (b) Interaction wrench estimation error
    ax2.plot(
        times,
        err_CF_hist,
        color="C1",
        label=r"$\|{}^{C}F_{\mathrm{int}}^{\mathrm{est}}-{}^{C}F_{\mathrm{true}}\|$",
    )
    ax2.set_title("(b) Interaction wrench estimation error")
    ax2.set_ylabel(r"Error norm $\|\cdot\|$")
    ax2.set_yscale("log")
    ax2.grid(True, linestyle="--", linewidth=0.5)
    ax2.legend(loc="lower right")

    # (c) Feedback torque rendering error
    ax3.plot(
        times,
        err_tau_fb_hist,
        color="C2",
        label=r"$\|\tau_{\mathrm{fb}}^{\mathrm{est}}-\tau_{\mathrm{fb}}^{\mathrm{true}}\|$",
    )
    ax3.set_title("(c) Haptic feedback torque error")
    ax3.set_xlabel(r"Time $t$ [s]")
    ax3.set_ylabel(r"Error norm $\|\cdot\|$")
    ax3.set_yscale("log")
    ax3.grid(True, linestyle="--", linewidth=0.5)
    ax3.legend(loc="lower right")

    fig_err.tight_layout()
    # fig_err.savefig("figs/haptic_errors_3subplots.png", dpi=300, bbox_inches="tight")

    # ------------------ 能量罐 + 缩放系数 ------------------
    fig_tank, ax1_t = plt.subplots(figsize=(6, 4))

    ax1_t.set_title("Energy tank dynamics and scaling factor")

    l1 = ax1_t.plot(
        times,
        E_hist,
        color="C0",
        label=r"Tank energy $E[k]$",
    )
    ax1_t.set_xlabel(r"Time $t$ [s]")
    ax1_t.set_ylabel(r"Energy $E$ [J]")
    ax1_t.grid(True, linestyle="--", linewidth=0.5)

    ax2_t = ax1_t.twinx()
    l2 = ax2_t.plot(
        times,
        alpha_smooth,
        color="C1",
        linestyle="--",
        label=r"Scaling $\alpha[k]$ (smoothed)",
    )
    ax2_t.set_ylabel(r"Scaling factor $\alpha$")

    # Legend 合并并右下角防遮挡
    lines = l1 + l2
    labels = [l.get_label() for l in lines]
    ax1_t.legend(lines, labels, loc="lower right")

    fig_tank.tight_layout()
    # fig_tank.savefig("figs/haptic_tank_smooth.png", dpi=300, bbox_inches="tight")

    plt.show()



if __name__ == "__main__":
    main()
