import os
import sys
import numpy as np
import pinocchio as pin
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
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


def safe_normalize(v, eps=1e-8):
    """Return v / ||v||, or None if the norm is too small."""
    n = np.linalg.norm(v)
    if n < eps:
        return None
    return v / n


def check_once(mapper, q_m, names, visualize=True):
    import numpy as np
    import open3d as o3d
    from utils.drwa import o3d_frame, o3d_lineset, o3d_sphere

    m_model, m_data = mapper.m_model, mapper.m_data
    s_model, s_data = mapper.s_model, mapper.s_data

    q_s, dbg = mapper.retargetting(q_m)

    # 主端 FK
    fk_update(m_model, m_data, q_m)
    pS_m, RS_m = pose_of_frame(m_model, m_data, names["m_shoulder"])
    pE_m, REf_m = pose_of_frame(m_model, m_data, names["m_elbow"])
    pW_m, RW_m = pose_of_frame(m_model, m_data, names["m_wrist"])
    p_m,  R_m  = pose_of_frame(m_model, m_data, names["m_ee"])

    # 从端 FK
    fk_update(s_model, s_data, q_s)
    pS_s, RS_s = pose_of_frame(s_model, s_data, names["s_shoulder"])
    pE_s, REf_s = pose_of_frame(s_model, s_data, names["s_elbow"])
    pW_s, RW_s = pose_of_frame(s_model, s_data, names["s_wrist"])
    p_s,  R_s  = pose_of_frame(s_model, s_data, names["s_ee"])

    # 误差
    pos_err = np.linalg.norm(p_m - p_s) * 1000.0  # mm
    rot_err = rot_error_deg(R_m, R_s)             # degree
    print(f"rot_err: {rot_err}, rot_err: {rot_err}")

    # 是否发生 reach 裁剪
    clipped = dbg["d_s"] < np.linalg.norm(
        m_data.oMf[m_model.getFrameId(names["m_wrist"])].translation
        - m_data.oMf[m_model.getFrameId(names["m_shoulder"])].translation) - 1e-9 \
            or dbg["d_s"] > mapper.L1 + mapper.L2 - mapper.eps_clip + 1e-9

    if visualize:
        geoms = []
        geoms.append(o3d_frame(np.eye(3), np.zeros(3), size=0.1))

        pts_m = np.vstack([pS_m, pE_m, pW_m, p_m])
        geoms.append(o3d_lineset(pts_m, color=(1.0, 0.0, 0.0)))
        geoms += [
            o3d_sphere(pS_m, 0.016, (1.0,0.4,0.4)),
            o3d_sphere(pE_m, 0.014, (1.0,0.3,0.3)),
            o3d_sphere(pW_m, 0.012, (1.0,0.2,0.2)),
            o3d_sphere(p_m,  0.012, (1.0,0.0,0.0)),
            o3d_frame(R_m, p_m, size=0.07)
        ]
        # 从臂（蓝）
        pts_s = np.vstack([pS_s, pE_s, pW_s, p_s])
        geoms.append(o3d_lineset(pts_s, color=(0.0, 0.4, 1.0)))
        geoms += [
            o3d_sphere(pS_s, 0.016, (0.4,0.6,1.0)),
            o3d_sphere(pE_s, 0.014, (0.3,0.5,1.0)),
            o3d_sphere(pW_s, 0.012, (0.2,0.4,1.0)),
            o3d_sphere(p_s,  0.012, (0.0,0.2,1.0)),
            o3d_frame(R_s, p_s, size=0.07)
        ]

        geoms.append(o3d_sphere(dbg["pW_s_ref"], 0.008, (0.2, 1.0, 0.2)))
        o3d.visualization.draw_geometries(geoms, window_name="SEW Retargeting (Master=Red, Slave=Blue)")
    return dict(pos_mm=pos_err, rot_deg=rot_err, clipped=clipped, dbg=dbg, q_s=q_s)


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 validation(mapper):
    q_m0 = pin.neutral(mapper.m_model)
    names = {
        "m_shoulder": "master_shoulder",
        "m_elbow": "master_forearm",
        "m_wrist": "master_wrist",
        "m_ee": "master_ee",
        "s_shoulder": "slave_shoulder",
        "s_elbow": "slave_forearm",
        "s_wrist": "slave_wrist",
        "s_ee": "slave_ee"
    }
    res0 = check_once(mapper, q_m0, names)
    print("[Neutral] pos_err = %.3f mm, rot_err = %.3f deg, clipped=%s" %
          (res0["pos_mm"], res0["rot_deg"], res0["clipped"]))

    # 2) 随机多组验证
    N = 100
    samples = random_qm(mapper, N=N, margin=0.1)
    pos_errs = []; rot_errs = []; n_clip = 0
    for q_m in samples:
        r = check_once(mapper, q_m, names)
        pos_errs.append(r["pos_mm"])
        rot_errs.append(r["rot_deg"])
        n_clip += int(r["clipped"])

    pos_errs = np.array(pos_errs)
    rot_errs = np.array(rot_errs)


# =======================================================
# 辅助函数：让 3D 轴等比例，图更“几何化”
# =======================================================
def set_equal_aspect_3d(ax, pts):
    """pts: (N,3) numpy array."""
    pts = np.asarray(pts)

    # 如果有 NaN 或 Inf，直接返回，不调坐标轴
    if not np.all(np.isfinite(pts)):
        return

    x_min, y_min, z_min = pts.min(axis=0)
    x_max, y_max, z_max = pts.max(axis=0)
    max_range = max(x_max - x_min, y_max - y_min, z_max - z_min)
    if max_range == 0:
        max_range = 1.0  # 防止除零

    x_mid = 0.5 * (x_max + x_min)
    y_mid = 0.5 * (y_max + y_min)
    z_mid = 0.5 * (z_max + z_min)
    r = 0.6 * max_range
    ax.set_xlim(x_mid - r, x_mid + r)
    ax.set_ylim(y_mid - r, y_mid + r)
    ax.set_zlim(z_mid - r, z_mid + r)



# =======================================================
#  绘图（论文插图专用）
# =======================================================

def plot_sew_geometry(pS_m, pE_m, pW_m,
                      pS_s, pE_s, pW_s_ref,
                      xhat_m, n_m, n_ref, phi_m):
    """
    Fig. SEW Geometry：
    - Master/Slave S-E-W
    - Shoulder-Wrist 方向 x̂_m
    - 上臂平面法向 n_m
    - 参考法向 n_ref
    """

    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111, projection='3d')

    # Master skeleton (粗实体＋虚线前臂)
    ax.plot([pS_m[0], pE_m[0]], [pS_m[1], pE_m[1]], [pS_m[2], pE_m[2]],
            'r-', lw=2.5, label='Master upper arm')
    ax.plot([pE_m[0], pW_m[0]], [pE_m[1], pW_m[1]], [pE_m[2], pW_m[2]],
            'r--', lw=2.0, label='Master forearm')

    # Slave skeleton (蓝色)
    ax.plot([pS_s[0], pE_s[0]], [pS_s[1], pE_s[1]], [pS_s[2], pE_s[2]],
            'b-', lw=2.5, label='Slave upper arm')
    ax.plot([pE_s[0], pW_s_ref[0]], [pE_s[1], pW_s_ref[1]], [pE_s[2], pW_s_ref[2]],
            'b--', lw=2.0, label='Slave forearm')

    # --- Master 上臂平面 patch（S_m, E_m, W_m） ---
    tri = np.vstack([pS_m, pE_m, pW_m])
    area = np.linalg.norm(np.cross(pE_m - pS_m, pW_m - pS_m))
    if area > 1e-8:
        # 使用 Poly3DCollection 画简单三角面片，避免 qhull
        verts = [tri]
        poly = Poly3DCollection(verts, alpha=0.08,
                            facecolor='r', edgecolor='none')
        ax.add_collection3d(poly)
        ax.plot_trisurf(tri[:, 0], tri[:, 1], tri[:, 2],
                    color='r', alpha=0.08, edgecolor='none')

    # Vectors at S_m
    arrow_len = 0.12
    ax.quiver(*pS_m, *(arrow_len * xhat_m),
              color='g', lw=2.0, normalize=False, label=r'$\hat{x}_m$')
    ax.quiver(*pS_m, *(arrow_len * n_m),
              color='m', lw=2.0, normalize=False, label=r'$n_m$')
    ax.quiver(*pS_m, *(arrow_len * n_ref),
              color='c', lw=2.0, normalize=False, label=r'$n_{\mathrm{ref}}$')

    # Labels
    ax.text(*pS_m, r'$S_m$', fontsize=10)
    ax.text(*pE_m, r'$E_m$', fontsize=10)
    ax.text(*pW_m, r'$W_m$', fontsize=10)
    ax.text(*pS_s, r'$S_s$', fontsize=10)
    ax.text(*pE_s, r'$E_s$', fontsize=10)
    ax.text(*pW_s_ref, r'$W_s$', fontsize=10)

    # 视角和轴
    pts_all = np.vstack([pS_m, pE_m, pW_m, pS_s, pE_s, pW_s_ref])
    set_equal_aspect_3d(ax, pts_all)
    ax.view_init(elev=22, azim=52)
    ax.set_xlabel("X"); ax.set_ylabel("Y"); ax.set_zlabel("Z")

    title_deg = phi_m * 180.0 / np.pi
    ax.set_title(r"SEW geometry (swivel $\phi_m = %.1f^\circ$)" % title_deg)
    ax.legend(loc='upper right', fontsize=8)
    plt.tight_layout()
    plt.show()


def plot_swivel_angle(pS_m, n_ref, n_m, xhat_m):
    """
    Fig. Swivel Angle:
    可视化 n_ref → n_m 绕 x̂_m 的旋转，并画出圆弧 φ_m。
    """

    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111, projection='3d')

    # 以 S_m 为圆心，在垂直于 xhat_m 的平面内画单位圆
    # 构造平面正交基 b1, b2
    tmp = np.array([1.0, 0.0, 0.0])
    if abs(np.dot(tmp, xhat_m)) > 0.9:
        tmp = np.array([0.0, 1.0, 0.0])
    b1 = tmp - np.dot(tmp, xhat_m) * xhat_m
    b1 = b1 / np.linalg.norm(b1)
    b2 = np.cross(xhat_m, b1)

    # 用 n_ref, n_m 在该平面上算出对应角度
    # 这里假设 n_ref 已经在平面内
    r = 0.12
    # Swivel angle
    cos_phi = np.dot(n_ref, n_m)
    sin_phi = np.dot(xhat_m, np.cross(n_ref, n_m))
    phi = np.arctan2(sin_phi, cos_phi)

    # 圆弧
    ts = np.linspace(0.0, phi, 80)
    arc_pts = []
    for t in ts:
        v = np.cos(t) * n_ref + np.sin(t) * np.cross(xhat_m, n_ref)
        v = v / np.linalg.norm(v)
        arc_pts.append(pS_m + r * v)
    arc_pts = np.asarray(arc_pts)
    ax.plot(arc_pts[:, 0], arc_pts[:, 1], arc_pts[:, 2], 'k-', lw=2.0, label=r'$\phi_m$')

    # 向量
    ax.quiver(*pS_m, *(r * n_ref), color='c', lw=2.0, normalize=False, label=r'$n_{\mathrm{ref}}$')
    ax.quiver(*pS_m, *(r * n_m),   color='m', lw=2.0, normalize=False, label=r'$n_m$')
    ax.quiver(*pS_m, *(r * xhat_m), color='g', lw=2.0, normalize=False, label=r'$\hat{x}_m$')

    ax.text(*pS_m, r'$S_m$', fontsize=10)

    pts_all = np.vstack([arc_pts, pS_m])
    set_equal_aspect_3d(ax, pts_all)
    ax.view_init(elev=20, azim=45)
    ax.set_xlabel("X"); ax.set_ylabel("Y"); ax.set_zlabel("Z")
    ax.set_title(r"Swivel angle definition")
    ax.legend(loc='upper right', fontsize=8)
    plt.tight_layout()
    plt.show()


def plot_two_sphere_sew(pS_s, pW_s_ref, L1, L2, pE_s):
    """
    Fig. Two-sphere intersection:
    - 球心 S_s, W_s，半径 L1, L2
    - 交点 E_s
    - 高亮三角形 S-E-W
    """

    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111, projection='3d')

    # 球面
    u, v = np.mgrid[0:2*np.pi:40j, 0:np.pi:20j]
    sphere1 = pS_s.reshape(3,1,1) + L1 * np.array([
        np.cos(u)*np.sin(v),
        np.sin(u)*np.sin(v),
        np.cos(v)
    ])
    sphere2 = pW_s_ref.reshape(3,1,1) + L2 * np.array([
        np.cos(u)*np.sin(v),
        np.sin(u)*np.sin(v),
        np.cos(v)
    ])
    ax.plot_surface(sphere1[0], sphere1[1], sphere1[2],
                    alpha=0.15, color='r', edgecolor='none')
    ax.plot_surface(sphere2[0], sphere2[1], sphere2[2],
                    alpha=0.15, color='b', edgecolor='none')

    # S-E-W 三角形
    tri = np.vstack([pS_s, pE_s, pW_s_ref])
    ax.plot_trisurf(tri[:, 0], tri[:, 1], tri[:, 2],
                    color='k', alpha=0.08, edgecolor='none')
    ax.plot([pS_s[0], pE_s[0], pW_s_ref[0], pS_s[0]],
            [pS_s[1], pE_s[1], pW_s_ref[1], pS_s[1]],
            [pS_s[2], pE_s[2], pW_s_ref[2], pS_s[2]],
            'k-', lw=2.0)

    # 关键点
    ax.scatter(*pE_s, color='k', s=50, label=r'$E_s$')
    ax.scatter(*pS_s, color='r', s=40, label=r'$S_s$')
    ax.scatter(*pW_s_ref, color='b', s=40, label=r'$W_s$')

    ax.set_xlabel("X"); ax.set_ylabel("Y"); ax.set_zlabel("Z")
    ax.set_title("Two-sphere intersection for SEW elbow")

    pts_all = np.vstack([pS_s, pE_s, pW_s_ref])
    set_equal_aspect_3d(ax, pts_all)
    ax.view_init(elev=20, azim=40)
    ax.legend(loc='upper right', fontsize=8)
    plt.tight_layout()
    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,
    )

    # 选一组姿态示例（这里用 neutral）
    # q_m = pin.neutral(mapper.m_model)
    samples = random_qm(mapper, N=500, margin=0.1)
    for q_m in samples: 
        q_s, dbg = mapper.retargetting(q_m)

        # 主端点
        fk_update(mapper.m_model, mapper.m_data, q_m)
        pS_m, _ = pose_of_frame(mapper.m_model, mapper.m_data, conf.m_shoulder_frame)
        pE_m, _ = pose_of_frame(mapper.m_model, mapper.m_data, conf.m_elbow_frame)
        pW_m, _ = pose_of_frame(mapper.m_model, mapper.m_data, conf.m_wrist_frame)

        # 从端（SEW 参考）
        pS_s = mapper.pS_s_fixed
        pE_s = dbg["pE_s"]
        pW_s_ref = dbg["pW_s_ref"]

        # 几何向量
        r_m = pW_m - pS_m
        xhat_m = safe_normalize(r_m)
        if xhat_m is None:
            raise RuntimeError("Shoulder–wrist vector is zero; check FK / frames.")

        # master 上臂平面法向
        n_m_raw = np.cross(pE_m - pS_m, pW_m - pS_m)
        n_m = safe_normalize(n_m_raw)

        # 参考法向（胸部 up 向量投影到垂直于 xhat_m 的平面）
        nref_tilde = mapper.up - (mapper.up @ xhat_m) * xhat_m
        n_ref = safe_normalize(nref_tilde)

        phi_m = dbg["phi_m"]

        degenerate_plane = (n_m is None or n_ref is None)
        if degenerate_plane:
            print("[Warn] S–E–W nearly collinear, swivel angle undefined for this posture.")
            # 用一个合理的默认法向，方便画 SEW overall 图；swivel 角度取 0
            if n_ref is None:
                # 如果连 n_ref 都 degenerate，就随便取一个与 xhat_m 垂直的向量
                tmp = np.array([1.0, 0.0, 0.0])
                if abs(np.dot(tmp, xhat_m)) > 0.9:
                    tmp = np.array([0.0, 1.0, 0.0])
                n_ref = safe_normalize(tmp - np.dot(tmp, xhat_m) * xhat_m)
            if n_m is None:
                n_m = n_ref.copy()
            phi_m = 0.0  # 退化情形下设成 0 仅用于示意

        # 绘制三类图
        print("Plotting SEW Geometry...")
        plot_sew_geometry(pS_m, pE_m, pW_m,
                        pS_s, pE_s, pW_s_ref,
                        xhat_m, n_m, n_ref, phi_m)

        if not degenerate_plane:
            print("Plotting Swivel Angle Geometry...")
            plot_swivel_angle(pS_m, n_ref, n_m, xhat_m)
        else:
            print("Skip swivel-angle figure for this degenerate posture (S–E–W nearly collinear).")

        print("Plotting Two-Sphere Intersection...")
        plot_two_sphere_sew(pS_s, pW_s_ref, mapper.L1, mapper.L2, pE_s)


if __name__ == '__main__':
    main()
