import numpy as np
import pinocchio as pin
from core.sew_mapper2 import SEWMapper, BallJointConfig


# -------------------------- quick test --------------------- ----- #

def test_once(mapper: SEWMapper) -> None:
    q_m7 = np.array([0.534, 0.314, -0.1, 2.14, 0.38, 0.38, -0.72], dtype=float)
    q_s, dbg = mapper.retarget(q_m7)

    # evaluate
    mapper._fk_master(_build_q7(mapper.m_model, mapper.m_qidx7, q_m7))
    mapper._fk_slave(q_s)
    Rm = mapper._rot(mapper.m_data, mapper.fid_mEE)
    Rs = mapper._rot(mapper.s_data, mapper.fid_sEE)
    ang_err = float(np.linalg.norm(pin.log3(Rs.T @ Rm)))
    pW = mapper._pos(mapper.s_data, mapper.fid_sW)
    pos_err = float(np.linalg.norm(pW - dbg["pW_s_ref"]))

    print("q_s(7):", np.array([q_s[i] for i in mapper.s_qidx7]))
    print("angle_err(rad):", ang_err)
    print("wrist_pos_err(m):", pos_err)

def _build_q7(model: pin.Model, qidx7, q7):
    q = pin.neutral(model)
    for i, idx in enumerate(qidx7):
        q[idx] = float(q7[i])
    return q


def test_fake():
    master_urdf = "../config/master_7dof.urdf"
    slave_urdf = "../config/slave_7dof.urdf"

    m_model = pin.buildModelFromUrdf(master_urdf)
    s_model = pin.buildModelFromUrdf(slave_urdf)

    # frames
    mS, mE, mW, mEE = "master_shoulder", "master_forearm", "master_wrist", "master_ee"
    sS, sE, sW, sEE = "slave_shoulder", "slave_forearm", "slave_wrist", "slave_ee"

    # joint order: [S1,S2,S3, EL, W1,W2,W3]
    master_joint_names = (
        "master_shoulder_pitch_joint", "master_shoulder_yaw_joint", "master_shoulder_roll_joint",
        "master_elbow_flex_joint",
        "master_wrist_roll_joint", "master_wrist_yaw_joint", "master_wrist_pitch_joint",
    )
    slave_joint_names = (
        "slave_shoulder_pitch_joint", "slave_shoulder_yaw_joint", "slave_shoulder_roll_joint",
        "slave_elbow_flex_joint",
        "slave_wrist_roll_joint", "slave_wrist_yaw_joint", "slave_wrist_pitch_joint",
    )

    # configs (your confirmed settings)
    s_sh_cfg = BallJointConfig(
        axis_order="yxy",
        joint_names=("slave_shoulder_pitch_joint", "slave_shoulder_yaw_joint", "slave_shoulder_roll_joint"),
        signs=(-1.0, 1.0, -1.0),
    )
    s_wr_cfg = BallJointConfig(
        axis_order="yzx",
        joint_names=("slave_wrist_roll_joint", "slave_wrist_yaw_joint", "slave_wrist_pitch_joint"),
        signs=(-1.0, 1.0, 1.0),
    )

    mapper = SEWMapper(
        master_model=m_model,
        slave_model=s_model,
        m_shoulder=mS, m_elbow=mE, m_wrist=mW, m_ee=mEE,
        s_shoulder=sS, s_elbow=sE, s_wrist=sW, s_ee=sEE,
        master_joint_names=master_joint_names,
        slave_joint_names=slave_joint_names,
        slave_shoulder_cfg=s_sh_cfg,
        slave_wrist_cfg=s_wr_cfg,
        slave_elbow_axis_local=np.array([1.0, 0.0, 0.0]),
        debug=False,
    )

    test_once(mapper)


def test_real():
    master_urdf = "../config/master_7dof.urdf"
    slave_urdf = "../config/real_slave_7dof.urdf"

    m_model = pin.buildModelFromUrdf(master_urdf)
    s_model = pin.buildModelFromUrdf(slave_urdf)

    # frames
    mS, mE, mW, mEE = "master_shoulder", "master_forearm", "master_wrist", "master_ee"
    sS, sE, sW, sEE = "R_SHOULDER_R_S", "R_ELBOW_R_S", "R_WRIST_R_S", "R_WRIST_R_S"

    # joint order: [S1,S2,S3, EL, W1,W2,W3]
    master_joint_names = (
        "master_shoulder_pitch_joint", "master_shoulder_yaw_joint", "master_shoulder_roll_joint",
        "master_elbow_flex_joint",
        "master_wrist_roll_joint", "master_wrist_yaw_joint", "master_wrist_pitch_joint",
    )
    slave_joint_names = (
        "R_SHOULDER_P", "R_SHOULDER_R", "R_SHOULDER_Y",
        "R_ELBOW_R",
        "R_WRIST_P", "R_WRIST_Y", "R_WRIST_R",
    )

    # configs (your confirmed settings)
    s_sh_cfg = BallJointConfig(
        axis_order="yxy",
        joint_names=("R_SHOULDER_P", "R_SHOULDER_R", "R_SHOULDER_Y"),
        signs=(-1.0, 1.0, -1.0),
    )
    s_wr_cfg = BallJointConfig(
        axis_order="yzx",
        joint_names=("R_WRIST_P", "R_WRIST_Y", "R_WRIST_R"),
        signs=(-1.0, 1.0, 1.0),
    )

    mapper = SEWMapper(
        master_model=m_model,
        slave_model=s_model,
        m_shoulder=mS, m_elbow=mE, m_wrist=mW, m_ee=mEE,
        s_shoulder=sS, s_elbow=sE, s_wrist=sW, s_ee=sEE,
        master_joint_names=master_joint_names,
        slave_joint_names=slave_joint_names,
        slave_shoulder_cfg=s_sh_cfg,
        slave_wrist_cfg=s_wr_cfg,
        slave_elbow_axis_local=np.array([1.0, 0.0, 0.0]),
        debug=False,
    )

    test_once(mapper)


if __name__ == "__main__":
    test_real()