/**
 ******************************************************************************
 * @file    arm.cpp
 * @brief   6-DOF robotic arm kinematics control implementation
 ******************************************************************************
 */

#include "arm6dof.h"
#include "toolbox/robotics.h"
#include "toolbox/matrix.h"
#include "user_math.h"
#include <string.h>

/* 静态变量 */
static robotics::Link* s_links = nullptr;
static robotics::Serial_Link<6>* s_robot = nullptr;
static bool s_initialized = false;

// 当前机械臂URDF(j1~j6)关节原始定义：parent->child 先 origin，再绕 joint axis 旋转。
// 这里显式按URDF链式连乘实现FK，避免将URDF joint origin强行映射到标准DH后产生几何失真。
static const float s_urdf_joint_origin_xyz[6][3] = {
    {0.0f,      0.0f,      0.024f},
    {-0.001395f,0.0f,      0.1015f},
    {0.3265f,   0.0f,     -0.0071975f},
    {0.0905f,   0.052775f, 0.0058025f},
    {0.001627f, 0.0f,      0.18467f},
    {0.10487f,  0.0013347f,0.0f},
};

// URDF顺序为(roll, pitch, yaw)。robotics::rpy2r输入顺序为(yaw, pitch, roll)。
static const float s_urdf_joint_origin_rpy[6][3] = {
    {0.0f,    0.0f,    0.0f},
    {1.5708f, 0.0f,   -1.5708f},
    {0.0f,    0.0f,    0.0f},
    {1.5708f, 0.0f,    3.1416f},
    {1.5708f, 0.0f,    0.0f},
    {1.5708f, 0.0f,    1.5708f},
};

// Link1~Link6 质心在各自link坐标系下的位置（直接取自 xin.urdf 的 inertial/origin）。
static const float s_urdf_link_com_xyz[6][3] = {
    {-0.0071819f,  0.00031034f,  0.061598f},
    { 0.22947f,   -0.0031632f,  -0.0022707f},
    { 0.086963f,   0.0016734f,  -0.017806f},
    { 0.000055295f,-0.002302704f,0.10972717f},
    { 0.056542f,   0.0013106f,   0.0010268f},
    {-0.0000077997f,0.000012767f,0.10159f},
};

static Matrixf<4, 4> Arm6dof_UrdfJointTransform(uint8_t index, float q) {
    const float roll  = s_urdf_joint_origin_rpy[index][0];
    const float pitch = s_urdf_joint_origin_rpy[index][1];
    const float yaw   = s_urdf_joint_origin_rpy[index][2];

    float ypr_origin[3] = {yaw, pitch, roll};
    Matrixf<3, 3> R_origin = robotics::rpy2r(Matrixf<3, 1>(ypr_origin));

    float p_raw[3] = {
        s_urdf_joint_origin_xyz[index][0],
        s_urdf_joint_origin_xyz[index][1],
        s_urdf_joint_origin_xyz[index][2],
    };
    Matrixf<3, 1> p_origin(p_raw);

    // 当前URDF关节轴均为 local-z：R_joint = RotZ(q)
    float ypr_joint[3] = {q, 0.0f, 0.0f};
    Matrixf<3, 3> R_joint = robotics::rpy2r(Matrixf<3, 1>(ypr_joint));

    return robotics::rp2t(R_origin, p_origin)
         * robotics::rp2t(R_joint, matrixf::zeros<3, 1>());
}

static Matrixf<4, 4> Arm6dof_UrdfForwardTransform(const Arm6dof_JointAngles_t* q,
                                                  uint8_t joint_num) {
    Matrixf<4, 4> T = matrixf::eye<4, 4>();
    for (uint8_t i = 0; i < joint_num; ++i) {
        T = T * Arm6dof_UrdfJointTransform(i, q->q[i]);
    }
    return T;
}

static float Arm6dof_UrdfPotentialEnergy(const Arm6dof_JointAngles_t* q) {
    constexpr float g = 9.81f;
    float U = 0.0f;

    Matrixf<4, 4> T_link = matrixf::eye<4, 4>();
    for (uint8_t i = 0; i < 6; ++i) {
        T_link = T_link * Arm6dof_UrdfJointTransform(i, q->q[i]);

        float p_com_raw[3] = {
            s_urdf_link_com_xyz[i][0],
            s_urdf_link_com_xyz[i][1],
            s_urdf_link_com_xyz[i][2],
        };
        Matrixf<3, 1> p_com_local(p_com_raw);
        Matrixf<4, 4> T_com = T_link * robotics::p2t(p_com_local);
        Matrixf<3, 1> p_com_world = robotics::t2p(T_com);

        U += s_links[i].m() * g * p_com_world[2][0];
    }

    return U;
}

static Matrixf<6, 1> Arm6dof_UrdfPoseError(const Matrixf<4, 4>& T_target,
                                           const Matrixf<4, 4>& T_cur) {
    Matrixf<6, 1> err;
    Matrixf<3, 1> pe = robotics::t2p(T_target) - robotics::t2p(T_cur);
    Matrixf<3, 1> we = robotics::t2twist(T_target * robotics::invT(T_cur)).block<3, 1>(3, 0);
    for (int k = 0; k < 3; ++k) {
        err[k][0] = pe[k][0];
        err[k + 3][0] = we[k][0];
    }
    return err;
}

static Matrixf<6, 6> Arm6dof_UrdfNumericJacobian(const Arm6dof_JointAngles_t* q) {
    Matrixf<6, 6> J = matrixf::zeros<6, 6>();
    const float eps = 1e-4f;

    for (int i = 0; i < 6; ++i) {
        Arm6dof_JointAngles_t q_plus = *q;
        Arm6dof_JointAngles_t q_minus = *q;
        q_plus.q[i] += eps;
        q_minus.q[i] -= eps;

        Matrixf<4, 4> T_plus = Arm6dof_UrdfForwardTransform(&q_plus, 6);
        Matrixf<4, 4> T_minus = Arm6dof_UrdfForwardTransform(&q_minus, 6);

        Matrixf<3, 1> p_plus = robotics::t2p(T_plus);
        Matrixf<3, 1> p_minus = robotics::t2p(T_minus);
        Matrixf<3, 1> dp = (p_plus - p_minus) / (2.0f * eps);

        Matrixf<3, 3> R_plus = robotics::t2r(T_plus);
        Matrixf<3, 3> R_minus = robotics::t2r(T_minus);
        Matrixf<3, 3> R_rel = R_plus * R_minus.trans();
        Matrixf<4, 1> angvec = robotics::r2angvec(R_rel);

        Matrixf<3, 1> w;
        w[0][0] = angvec[0][0] * angvec[3][0] / (2.0f * eps);
        w[1][0] = angvec[1][0] * angvec[3][0] / (2.0f * eps);
        w[2][0] = angvec[2][0] * angvec[3][0] / (2.0f * eps);

        J[0][i] = dp[0][0];
        J[1][i] = dp[1][0];
        J[2][i] = dp[2][0];
        J[3][i] = w[0][0];
        J[4][i] = w[1][0];
        J[5][i] = w[2][0];
    }

    return J;
}

/**
 * @brief 初始化机械臂运动学模型
 */
void Arm6dof_Init(const Arm6dof_DHParams_t dh_params[6]) {
    if (dh_params == nullptr) {
        return;
    }
    
    // 创建DH链表
    if (s_links != nullptr) {
        delete[] s_links;
    }
    s_links = new robotics::Link[6];
    
    for (int i = 0; i < 6; i++) {
        // 构建质心向量（DH连杆坐标系下，由URDF经Rz(-theta_offset)旋转得到）
        Matrixf<3, 1> rc_vec;
        rc_vec[0][0] = dh_params[i].rc[0];
        rc_vec[1][0] = dh_params[i].rc[1];
        rc_vec[2][0] = dh_params[i].rc[2];

        s_links[i] = robotics::Link(
            dh_params[i].theta,         // 初始theta值（通常为0）
            dh_params[i].d,             // 连杆偏移 (m)
            dh_params[i].a,             // 连杆长度 (m)
            dh_params[i].alpha,         // 扭转角 (rad)
            robotics::R,                // 旋转关节
            dh_params[i].theta_offset,  // theta偏移 (rad)
            dh_params[i].qmin,          // 最小角度 (rad)
            dh_params[i].qmax,          // 最大角度 (rad)
            dh_params[i].m,             // 质量 (kg)
            rc_vec                      // 质心坐标（DH连杆坐标系，m）
        );
    }
    
    // 创建机器人模型
    if (s_robot != nullptr) {
        delete s_robot;
    }
    s_robot = new robotics::Serial_Link<6>(s_links);
    
    s_initialized = true;
}

/**
 * @brief 正运动学
 */
int Arm6dof_ForwardKinematics(const Arm6dof_JointAngles_t* q, Arm6dof_Pose_t* pose) {
    if (!s_initialized || q == nullptr || pose == nullptr) {
        return -1;
    }

    // 使用URDF原生链式FK：T = Π( T_origin_i * RotZ(q_i) )
    Matrixf<4, 4> T = Arm6dof_UrdfForwardTransform(q, 6);
    
    // 提取位置 (x, y, z)
    Matrixf<3, 1> pos = robotics::t2p(T);
    pose->x = pos[0][0];
    pose->y = pos[1][0];
    pose->z = pos[2][0];
    
    // 提取姿态 (RPY欧拉角)
    Matrixf<3, 1> rpy = robotics::t2rpy(T);
    pose->yaw   = rpy[0][0];
    pose->pitch = rpy[1][0];
    pose->roll  = rpy[2][0];
    
    return 0;
}

/**
 * @brief 获取指定关节的位姿
 */
int Arm6dof_GetJointPose(const Arm6dof_JointAngles_t* q, uint8_t joint_num, Arm6dof_Pose_t* pose) {
    if (!s_initialized || q == nullptr || pose == nullptr) {
        return -1;
    }
    
    if (joint_num < 1 || joint_num > 6) {
        return -1;
    }
    
    // 计算到第k个关节的URDF链式变换矩阵
    Matrixf<4, 4> T = Arm6dof_UrdfForwardTransform(q, joint_num);
    
    // 提取位置和姿态
    Matrixf<3, 1> pos = robotics::t2p(T);
    pose->x = pos[0][0];
    pose->y = pos[1][0];
    pose->z = pos[2][0];
    
    Matrixf<3, 1> rpy = robotics::t2rpy(T);
    pose->yaw   = rpy[0][0];
    pose->pitch = rpy[1][0];
    pose->roll  = rpy[2][0];
    
    return 0;
}

/**
 * @brief 逆运动学
 */
int Arm6dof_InverseKinematics(const Arm6dof_Pose_t* pose, const Arm6dof_JointAngles_t* q_init, 
                          Arm6dof_JointAngles_t* q_result, float tol, uint16_t max_iter) {
    if (!s_initialized || pose == nullptr || q_result == nullptr) {
        return -1;
    }
    
    // 构造目标变换矩阵
    Matrixf<3, 1> rpy_vec;
    rpy_vec[0][0] = pose->yaw;
    rpy_vec[1][0] = pose->pitch;
    rpy_vec[2][0] = pose->roll;
    Matrixf<3, 3> R_target = robotics::rpy2r(rpy_vec);
    
    Matrixf<3, 1> p_target;
    p_target[0][0] = pose->x;
    p_target[1][0] = pose->y;
    p_target[2][0] = pose->z;
    
    Matrixf<4, 4> T_target = robotics::rp2t(R_target, p_target);
    
    // 初始猜测值
    Matrixf<6, 1> q_guess = matrixf::zeros<6, 1>();
    if (q_init != nullptr) {
        for (int i = 0; i < 6; i++) {
            q_guess[i][0] = q_init->q[i];
        }
    }

    // 基于URDF链式FK的LM数值迭代，避免FK/IK模型不一致。
    Matrixf<6, 1> q_cur = q_guess;
    float lambda = 1.0f;
    constexpr float lambda_up = 10.0f;
    constexpr float lambda_down = 0.1f;
    constexpr float lambda_max = 1e6f;
    constexpr float lambda_min = 1e-6f;

    for (uint16_t iter = 0; iter < max_iter; ++iter) {
        Arm6dof_JointAngles_t q_cur_struct;
        for (int i = 0; i < 6; ++i) {
            q_cur_struct.q[i] = q_cur[i][0];
        }

        Matrixf<4, 4> T_cur = Arm6dof_UrdfForwardTransform(&q_cur_struct, 6);
        Matrixf<6, 1> err = Arm6dof_UrdfPoseError(T_target, T_cur);
        float err_norm = err.norm();
        if (err_norm < tol) {
            break;
        }

        Matrixf<6, 6> J = Arm6dof_UrdfNumericJacobian(&q_cur_struct);
        Matrixf<6, 6> H = J.trans() * J + lambda * matrixf::eye<6, 6>();
        Matrixf<6, 1> g = J.trans() * err;
        Matrixf<6, 1> dq = matrixf::inv(H) * g;

        if (dq[0][0] == INFINITY || dq[0][0] != dq[0][0]) {
            lambda *= lambda_up;
            if (lambda > lambda_max) {
                return -1;
            }
            continue;
        }

        Matrixf<6, 1> q_new = q_cur + dq;

        // 对循环关节折叠，避免等效角导致数值发散。
        for (int i = 0; i < 6; ++i) {
            constexpr float TWO_PI = 2.0f * (float)M_PI;
            bool is_cyclic = (s_links[i].qmin_ < 0.1f) && (s_links[i].qmax_ > TWO_PI - 0.1f);
            if (is_cyclic) {
                q_new[i][0] = math::loopLimit(q_new[i][0], s_links[i].qmin_, s_links[i].qmax_);
            }
        }

        Arm6dof_JointAngles_t q_new_struct;
        for (int i = 0; i < 6; ++i) {
            q_new_struct.q[i] = q_new[i][0];
        }

        Matrixf<4, 4> T_new = Arm6dof_UrdfForwardTransform(&q_new_struct, 6);
        Matrixf<6, 1> new_err = Arm6dof_UrdfPoseError(T_target, T_new);
        float new_err_norm = new_err.norm();

        if (new_err_norm < err_norm) {
            q_cur = q_new;
            lambda *= lambda_down;
            if (lambda < lambda_min) {
                lambda = lambda_min;
            }
        } else {
            lambda *= lambda_up;
            if (lambda > lambda_max) {
                return -1;
            }
        }
    }

    for (int i = 0; i < 6; ++i) {
        q_result->q[i] = q_cur[i][0];
    }

    // 验证解的精度：位置误差 + 姿态误差双重校验
    Arm6dof_JointAngles_t q_verify_struct;
    for (int i = 0; i < 6; ++i) {
        q_verify_struct.q[i] = q_cur[i][0];
    }
    Matrixf<4, 4> T_verify = Arm6dof_UrdfForwardTransform(&q_verify_struct, 6);
    Matrixf<3, 1> p_verify = robotics::t2p(T_verify);
    
    float pos_error = (p_verify - p_target).norm();
    
    if (pos_error > tol * 10.0f) {  // 位置误差过大，未收敛
        return -1;
    }
    
    // 姿态误差：计算两旋转矩阵的测地距离 theta = arccos((trace(R_v * R_t^T) - 1) / 2)
    // 使用旋转矩阵方法避免 RPY 奇异点和角度折叠问题
    Matrixf<3, 3> R_verify = robotics::t2r(T_verify);
    Matrixf<3, 3> R_diff   = R_verify * R_target.trans();
    float trace_val = R_diff[0][0] + R_diff[1][1] + R_diff[2][2];
    float cos_theta = (trace_val - 1.0f) * 0.5f;
    // 数值钳位防止 acosf 域溢出
    if (cos_theta >  1.0f) cos_theta =  1.0f;
    if (cos_theta < -1.0f) cos_theta = -1.0f;
    float ang_error = acosf(cos_theta);  // [0, π]
    
    const float ang_tol = 0.1f;  // 姿态收敛阈值 (rad ≈ 5.7°)
    if (ang_error > ang_tol) {   // 姿态未收敛（收敛到错误分支或近奇异点）
        return -1;
    }
    
    return 0;
}

/**
 * @brief 计算重力补偿力矩
 * @details 与当前URDF链式运动学保持一致：通过势能 U(q)=Σ(mgz) 的数值梯度
 *          τg = ∂U/∂q 计算重力补偿力矩。
 */
int Arm6dof_GravityCompensation(const Arm6dof_JointAngles_t* q, float gravity_torques[6]) {
    if (!s_initialized || q == nullptr || gravity_torques == nullptr) {
        return -1;
    }
    
    // // 关节角度
    // Matrixf<6, 1> q_mat;
    // for (int i = 0; i < 6; i++) {
    //     q_mat[i][0] = q->q[i];
    // }
    
    // // 速度和加速度均为零 —— 只计算重力项
    // Matrixf<6, 1> qv = matrixf::zeros<6, 1>();
    // Matrixf<6, 1> qa = matrixf::zeros<6, 1>();
    
    // // 调用牛顿-欧拉逆动力学
    // Matrixf<6, 1> torques = s_robot->rne(q_mat, qv, qa);
    
    // for (int i = 0; i < 6; i++) {
    //     gravity_torques[i] = torques[i][0];
    // }

        constexpr float eps = 1e-4f;
    for (int i = 0; i < 6; ++i) {
        Arm6dof_JointAngles_t q_plus = *q;
        Arm6dof_JointAngles_t q_minus = *q;
        q_plus.q[i] += eps;
        q_minus.q[i] -= eps;

        float U_plus = Arm6dof_UrdfPotentialEnergy(&q_plus);
        float U_minus = Arm6dof_UrdfPotentialEnergy(&q_minus);
        gravity_torques[i] = (U_plus - U_minus) / (2.0f * eps);
    }
    
    return 0;
}