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

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from core.sew_mapper import SEWMapper, fk_update, pose_of_frame, rot_error_deg

# ...（这里保留你之前的 safe_normalize、random_qm、几何绘图函数等）...
def random_qm(mapper, N=50, margin=0.2):
    lb = []; ub = []
    for j in mapper.m_model.joints[1:]:  # skip universe
        if j.nq == 1:  # 1-DoF revolute
            jid = j.id
            iq = j.idx_q
            lb.append(mapper.m_model.lowerPositionLimit[iq])
            ub.append(mapper.m_model.upperPositionLimit[iq])
    lb = np.array(lb); ub = np.array(ub)
    rng = (ub - lb)
    lb2 = lb + margin * rng
    ub2 = ub - margin * rng

    qs = []
    for _ in range(N):
        v = lb2 + np.random.rand(len(lb2)) * (ub2 - lb2)
        q_full = pin.neutral(mapper.m_model)
        k = 0
        for j in mapper.m_model.joints[1:]:
            if j.nq == 1:
                q_full[j.idx_q] = v[k]; k += 1
        qs.append(q_full)
    return qs



def evaluate_sew_statistics(mapper, conf, N=500, margin=0.1, save_dir="assets"):
    """
    统计验证 SEW retargeting 的精度，并输出论文用图：
      (1) 腕点位置误差直方图
      (2) 腕点姿态误差直方图
      (3) （可选）位置误差 vs 肩-腕距离散点
    """
    os.makedirs(save_dir, exist_ok=True)

    # 随机采样 master 关节
    qs_m = random_qm(mapper, N=N, margin=margin)

    pos_errs = []    # mm
    rot_errs = []    # deg
    dists_sw = []    # shoulder–wrist distance (m)
    clipped_flags = []

    for q_m in qs_m:
        # SEW 映射
        q_s, dbg = mapper.retargetting(q_m)

        # master FK
        pS_m, _   = pose_of_frame(mapper.m_model, mapper.m_data, conf.m_shoulder_frame)
        pE_m, RE_m = pose_of_frame(mapper.m_model, mapper.m_data, conf.m_ee_frame)

        # slave FK
        fk_update(mapper.s_model, mapper.s_data, q_s)
        pE_s, RE_s = pose_of_frame(mapper.s_model, mapper.s_data, conf.s_ee_frame)

        # 误差（改成 EE）
        e_p = np.linalg.norm(pE_s - pE_m) * 1000.0   # 末端位置误差 [mm]
        e_R = rot_error_deg(RE_m, RE_s)             # 末端姿态误差 [deg]

        pos_errs.append(e_p)
        rot_errs.append(e_R)

        d_m = np.linalg.norm(pE_m - pS_m)                # m
        dists_sw.append(d_m)

        # 是否发生 reach 裁剪（根据 d_s vs master 原始 d_m、以及 L1+L2 区间）
        d_s = dbg["d_s"]
        clipped = (d_s < d_m - 1e-9) or (d_s > mapper.L1 + mapper.L2 - mapper.eps_clip + 1e-9)
        clipped_flags.append(clipped)

    pos_errs = np.asarray(pos_errs)
    rot_errs = np.asarray(rot_errs)
    dists_sw = np.asarray(dists_sw)
    clipped_flags = np.asarray(clipped_flags, dtype=bool)

    # --------- 打印统计量 (你可以把结果手工抄进论文表格) ----------
    def summary_str(x):
        return f"mean={x.mean():.3f}, std={x.std():.3f}, max={x.max():.3f}"

    print("=== SEW retargeting statistics over {} samples ===".format(N))
    print("Wrist position error [mm]:", summary_str(pos_errs))
    print("Wrist orientation error [deg]:", summary_str(rot_errs))
    print("Clipped configurations: {} / {} ({:.1f}%)"
          .format(clipped_flags.sum(), N, 100.0 * clipped_flags.mean()))

    # --------- 图 1: 位置 & 姿态误差直方图 ----------
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))

    axes[0].hist(pos_errs, bins=30, color='C0', alpha=0.8)
    axes[0].set_xlabel("Wrist position error $e_p$ [mm]")
    axes[0].set_ylabel("Count")
    axes[0].set_title("(a) Distribution of $e_p$")

    axes[1].hist(rot_errs, bins=30, color='C1', alpha=0.8)
    axes[1].set_xlabel("Wrist orientation error $e_R$ [deg]")
    axes[1].set_ylabel("Count")
    axes[1].set_title("(b) Distribution of $e_R$")

    plt.tight_layout()
    fig.savefig(os.path.join(save_dir, "sew_error_hist.png"), dpi=300)
    # plt.show()

    # --------- 图 2: 位置误差 vs 肩-腕距离 (可视化剪裁区) ----------
    fig2, ax2 = plt.subplots(figsize=(5, 4))
    ax2.scatter(dists_sw[~clipped_flags], pos_errs[~clipped_flags],
                s=15, c='C0', label="Unclipped")
    ax2.scatter(dists_sw[clipped_flags], pos_errs[clipped_flags],
                s=25, c='C3', marker='x', label="Clipped")
    ax2.set_xlabel("Shoulder–wrist distance $d_m$ [m]")
    ax2.set_ylabel("Position error $e_p$ [mm]")
    ax2.set_title("SEW error vs. reach distance")
    ax2.grid(True, linestyle='--', linewidth=0.5)
    ax2.legend(loc="upper left")

    fig2.tight_layout()
    fig2.savefig(os.path.join(save_dir, "sew_error_scatter.png"), dpi=300)
    plt.show()


def main():
    conf = OmegaConf.load("./config/config.yaml")

    master_model, _, _ = pin.buildModelsFromUrdf(str(conf.master_urdf))
    slave_model, _, _  = pin.buildModelsFromUrdf(str(conf.slave_urdf))

    mapper = SEWMapper(
        master_model=master_model,
        slave_model=slave_model,
        m_shoulder_frame=conf.m_shoulder_frame,
        m_elbow_frame=conf.m_elbow_frame,
        m_wrist_frame=conf.m_wrist_frame,
        m_ee_frame=conf.m_ee_frame,
        s_shoulder_frame=conf.s_shoulder_frame,
        s_elbow_frame=conf.s_elbow_frame,
        s_wrist_frame=conf.s_wrist_frame,
        s_ee_frame=conf.s_ee_frame,
        slave_joint_names=conf.sew_mapper.slave_joint_names,
        up_dir=np.array(conf.sew_mapper.up_dir),
        eps_clip=conf.sew_mapper.eps_clip,
    )

    # 1) 如果要画几何示意图（方法部分），可以单独写一个函数；这里我们只跑统计
    evaluate_sew_statistics(mapper, conf, N=500, margin=0.1, save_dir="assets")


if __name__ == "__main__":
    main()