# -*- coding: utf-8 -*-
import numpy as np
import pinocchio as pin
from pathlib import Path


def hat(v):
    x, y, z = v
    return np.array([[0, -z, y],
                     [z,  0,-x],
                     [-y, x, 0]], dtype=float)


def normalize(v, eps=1e-12):
    n = np.linalg.norm(v)
    if n < eps:
        return v * 0.0
    return v / n


def rodrigues(u, phi):
    u = normalize(u)
    K = hat(u)
    return np.eye(3) + np.sin(phi)*K + (1-np.cos(phi))*(K@K)


def minimal_rotation_align(a, b):
    a = normalize(a); b = normalize(b)
    v = np.cross(a, b)
    s = np.linalg.norm(v)
    c = float(np.dot(a, b))
    if s < 1e-12:
        # parallel or anti-parallel
        if c > 0.0:
            return np.eye(3)
        # 180°: choose any axis orthogonal to a
        axis = normalize(np.array([1.0,0,0]) if abs(a[0])<0.9 else np.array([0,1.0,0]))
        axis = normalize(np.cross(a, axis))
        K = hat(axis)
        return np.eye(3) + 2*(K@K)  # R = I + 2[K]^2 for 180°
    K = hat(v/s)
    return np.eye(3) + K*s + (1-c)/(s*s) * (K@K)


def euler_zyx_from_R(R):
    """Return Z-Y-X Euler angles (about z,y,x of the shoulder frame)."""
    sy = -R[2,0]
    cy = np.sqrt(max(0.0, 1.0 - sy*sy))
    if cy > 1e-9:
        z = np.arctan2(R[1,0], R[0,0])
        y = np.arctan2(sy, cy)
        x = np.arctan2(R[2,1], R[2,2])
    else:
        # Gimbal case: cy ~ 0
        z = np.arctan2(-R[0,1], R[1,1])
        y = np.arctan2(sy, cy)
        x = 0.0
    return np.array([z, y, x], dtype=float)

def euler_xzy_from_R(R):
    """
    Decompose R ≈ Rx(a) * Rz(b) * Ry(c)
    Return (a, b, c)
    """
    # from derivation:
    # sb = -R[0,1]; cb = sqrt(R[0,0]^2 + R[0,2]^2)
    sb = -R[0,1]
    cb = np.sqrt(max(0.0, R[0,0]**2 + R[0,2]**2))
    b  = np.arctan2(sb, cb)
    # c = atan2(R[0,2], R[0,0])
    c  = np.arctan2(R[0,2], R[0,0])
    # a from R[1,1] = ca*cb and R[2,1] = sa*cb
    if cb < 1e-9:
        # singular: cb≈0 => b≈±pi/2，退化时把 a=0，c 吸收残差（简化处理）
        a = 0.0
    else:
        ca = np.clip(R[1,1] / cb, -1.0, 1.0)
        sa = np.clip(R[2,1] / cb, -1.0, 1.0)
        a  = np.arctan2(sa, ca)
    return np.array([a, b, c], dtype=float)

def euler_zy_from_R(R):
    """Return Z-Y Euler angles for wrist (about z then y)."""
    # R ≈ Rz(z) * Ry(y)
    # y = asin(R[0,2])? We'll use standard decomposition.
    # From R = Rz*Ry:
    # R[2,0] = -sin(y)
    y = np.arcsin(np.clip(R[0,2], -1.0, 1.0))
    cy = np.cos(y)
    if abs(cy) < 1e-9:
        z = 0.0
    else:
        z = np.arctan2(-R[0,1]/cy, R[0,0]/cy)
    return np.array([z, y], dtype=float)


def rot_error_deg(RA, RB):
    # 旋转误差：以李群对数映射的范数（弧度），再转度
    R = RA.T @ RB
    w = pin.log3(R)
    return np.linalg.norm(w) * 180.0/np.pi

def pose_of_frame(model, data, frame_name):
    fid = model.getFrameId(frame_name)
    oMf = data.oMf[fid]
    return oMf.translation.copy(), oMf.rotation.copy()

def fk_update(model, data, q):
    pin.forwardKinematics(model, data, q)
    pin.updateFramePlacements(model, data)



class SEWMapper:
    """
    SEW Mapper (closed-form, no numeric IK):
    - Loads master & slave URDFs
    - Computes heteromorphic retargeting from master EE pose to slave q_s using SEW geometry
    Assumptions:
      * Slave arm is 7-DoF in the order: 3 shoulder + 1 elbow + 3 wrist (axes orthogonal at shoulder/wrist frames).
      * Provide correct link/frame names and joint name order for your model.
    """

    def __init__(
        self,
        master_model: pin.Model,
        slave_model: pin.Model,
        # Frame/link names
        m_shoulder_frame: str,
        m_elbow_frame: str,
        m_wrist_frame: str,
        m_ee_frame: str,
        s_shoulder_frame: str,
        s_elbow_frame: str,
        s_wrist_frame: str,
        s_ee_frame: str,
        # Joint name order for slave 7-DoF: [S1,S2,S3, EL, W1, W2, W3]
        slave_joint_names: list,
        # world/up direction and safety margins
        up_dir=np.array([0,0,1.0]),
        eps_clip=1e-3,
    ):
        # Load master
        self.m_model = master_model
        self.m_data  = self.m_model.createData()

        # Load slave
        self.s_model = slave_model
        self.s_data  = self.s_model.createData()

        # Frame IDs
        self.fid_mS = self._get_frame_id(self.m_model, m_shoulder_frame)
        self.fid_mE = self._get_frame_id(self.m_model, m_elbow_frame)
        self.fid_mW = self._get_frame_id(self.m_model, m_wrist_frame)
        self.fid_mEE= self._get_frame_id(self.m_model, m_ee_frame)

        self.fid_sS = self._get_frame_id(self.s_model, s_shoulder_frame)
        self.fid_sE = self._get_frame_id(self.s_model, s_elbow_frame)
        self.fid_sW = self._get_frame_id(self.s_model, s_wrist_frame)
        self.fid_sEE= self._get_frame_id(self.s_model, s_ee_frame)

        # Up direction and epsilon
        self.up = normalize(up_dir)
        self.eps_clip = float(eps_clip)

        # Query slave joint indices (order critical)
        self.s_joint_ids = [self.s_model.getJointId(n) for n in slave_joint_names]
        self.s_qidx      = [self.s_model.joints[jid].idx_q for jid in self.s_joint_ids]
        assert len(self.s_qidx) == 7, "Provide 7 slave joints in order [S1,S2,S3, EL, W1, W2, W3]"

        # Pre-compute slave segment lengths L1, L2 in reference (zero) config
        self.qs0 = pin.neutral(self.s_model)
        self._update_fk_slave(self.qs0)
        pS = self._frame_pos(self.s_data, self.fid_sS)
        pE = self._frame_pos(self.s_data, self.fid_sE)
        pW = self._frame_pos(self.s_data, self.fid_sW)
        self.L1 = float(np.linalg.norm(pE - pS))
        self.L2 = float(np.linalg.norm(pW - pE))

        # Fixed slave shoulder location in world
        self.pS_s_fixed = pS.copy()

    # ------------------ Utilities ------------------ #
    def _get_frame_id(self, model, name):
        try:
            return model.getFrameId(name)
        except:
            # fallback: also try link->frame mapping via joint name
            return model.getFrameId(name)

    def _update_fk_master(self, q_m):
        pin.forwardKinematics(self.m_model, self.m_data, q_m)
        pin.updateFramePlacements(self.m_model, self.m_data)

    def _update_fk_slave(self, q_s):
        pin.forwardKinematics(self.s_model, self.s_data, q_s)
        pin.updateFramePlacements(self.s_model, self.s_data)

    def _frame_pos(self, data, fid):
        return data.oMf[fid].translation.copy()

    def _frame_rot(self, data, fid):
        return data.oMf[fid].rotation.copy()

    # ------------------ Core Retarget ------------------ #
    def retargetting(self, q_m, q_s_init=None):
        """
        Input:
          q_m     : master joint configuration (np.ndarray, size = master nq)
          q_s_init: optional slave initial seed (ignored for geometry, used only to keep continuity of angle unwrap if desired)
        Output:
          q_s     : slave joint angles (np.ndarray, size = slave nq)
          debug   : dict with intermediate targets (pE_s, pW_s_ref, phi_m, theta)
        """
        # 1) Master FK and SEW quantities
        self._update_fk_master(q_m)
        pS_m = self._frame_pos(self.m_data, self.fid_mS)
        pE_m = self._frame_pos(self.m_data, self.fid_mE)
        pW_m = self._frame_pos(self.m_data, self.fid_mW)
        # RW_m = self._frame_rot(self.m_data, self.fid_mW)
        RE_m = self._frame_rot(self.m_data, self.fid_mEE) 

        r_m = pW_m - pS_m
        d_m = float(np.linalg.norm(r_m))
        xhat_m = normalize(r_m)

        # Swivel (master)
        nm_raw = np.cross(pE_m - pS_m, pW_m - pS_m)
        nm = normalize(nm_raw)
        nref_tilde = self.up - np.dot(self.up, xhat_m) * xhat_m
        if np.linalg.norm(nref_tilde) < 1e-6:
            ey = np.array([0,1.0,0])
            nref_tilde = ey - np.dot(ey, xhat_m) * xhat_m
        nref = normalize(nref_tilde)
        num = np.dot(xhat_m, np.cross(nref, nm))
        den = float(np.dot(nref, nm))
        phi_m = np.arctan2(num, den)

        # 2) Wrist pose and reach clipping on slave
        d_min = abs(self.L1 - self.L2) + self.eps_clip
        d_max = (self.L1 + self.L2) - self.eps_clip
        d_s = np.clip(d_m, d_min, d_max)
        xhat_s = xhat_m.copy()
        pS_s = self.pS_s_fixed
        pW_s_ref = pS_s + d_s * xhat_s
        R_ee_ref = RE_m.copy()   # preserve orientation

        # 3) Two-sphere elbow construction on slave
        d = float(np.linalg.norm(pW_s_ref - pS_s))
        e3 = rodrigues(xhat_s, phi_m) @ nref
        e3 = normalize(e3)
        e2 = normalize(np.cross(xhat_s, e3))
        cos_th = (self.L1**2 + d**2 - self.L2**2) / (2*self.L1*d)
        cos_th = float(np.clip(cos_th, -1.0, 1.0))
        th = np.arccos(cos_th)
        sin_th = np.sqrt(max(0.0, 1.0 - cos_th*cos_th))
        pE_s = pS_s + self.L1*(cos_th * xhat_s + sin_th * e2)

        # 4) Joint reconstruction (assumes 3-1-3 structure with Z-Y-X at shoulder and Z-Y at wrist)
        q_s = pin.neutral(self.s_model) if q_s_init is None else q_s_init.copy()

        # (a) Set shoulder (first 3 joints): align upper-arm vector
        u = normalize(pE_s - pS_s)
        y_axis = e3               # 让“肘轴=局部Y”严格对齐 SEW 的平面法向
        x_axis = u
        z_axis = normalize(np.cross(x_axis, y_axis))
        RS_des = np.column_stack([x_axis, y_axis, z_axis])


        dz, dy, dx = euler_zyx_from_R(RS_des)               # R ≈ Rz(dz)*Ry(dy)*Rx(dx)
        q_s[self.s_qidx[0]] = dz
        q_s[self.s_qidx[1]] = dy
        q_s[self.s_qidx[2]] = dx
        # (b) Elbow flex = pi - theta
        fhat = normalize(pW_s_ref - pE_s)
        u = normalize(pE_s - pS_s)# 上臂方向（从肩指向肘）
        q_elbow = np.arctan2(np.dot(e3, np.cross(u, fhat)), np.dot(u, fhat))# 计算把 u 绕 e3 旋到 fhat 的有符号角：atan2( 轴·(u×f),  u·f )
        q_s[self.s_qidx[3]] = q_elbow
        self._update_fk_slave(q_s)
        # (c) Wrist orientation match with X–Z–Y on wrist frame
        RW_cur = self._frame_rot(self.s_data, self.fid_sW)
        R_needed = RW_cur.T @ R_ee_ref         
        a_b_c = euler_xzy_from_R(R_needed)      # (roll-X, yaw-Z, pitch-Y)
        q_s[self.s_qidx[4]] = a_b_c[0]
        q_s[self.s_qidx[5]] = a_b_c[1]
        q_s[self.s_qidx[6]] = a_b_c[2]

        debug = dict(
            phi_m=phi_m,
            d_s=d_s,
            pW_s_ref=pW_s_ref,
            pE_s=pE_s,
            theta=th
        )
        return q_s, debug