#include "robot_kinematics.hpp"
#include <cmath>
#include <stdexcept>
#include <iostream>

namespace robot_model
{

RobotKinematics::RobotKinematics(std::shared_ptr<RobotModel> robot_model)
    : robot_model_(robot_model)
{
    if (!robot_model_)
    {
        throw std::invalid_argument("Robot model pointer cannot be null");
    }
}

Eigen::Matrix3d RobotKinematics::rotate_z(double angle)
{
    Eigen::Matrix3d mat;
    mat << std::cos(angle), -std::sin(angle), 0,
           std::sin(angle),  std::cos(angle), 0,
           0,                0,               1;
    return mat;
}

Eigen::Matrix3d RobotKinematics::rotate_y(double angle)
{
    Eigen::Matrix3d mat;
    mat << std::cos(angle),  0, std::sin(angle),
           0,                1, 0,
           -std::sin(angle), 0, std::cos(angle);
    return mat;
}

Eigen::Matrix3d RobotKinematics::rotate_x(double angle)
{
    Eigen::Matrix3d mat;
    mat << 1, 0,                 0,
           0, std::cos(angle), -std::sin(angle),
           0, std::sin(angle),  std::cos(angle);
    return mat;
}

Eigen::Vector3d RobotKinematics::arm_forward_kinematics(const std::vector<double>& joint_angles)
{
    const auto& arm = robot_model_->get_arm();
    
    // 检查关节角度数量是否正确
    if (joint_angles.size() != arm.joints.size())
    {
        throw std::invalid_argument("Invalid number of joint angles for arm");
    }
    
    // 连杆长度
    double l1 = arm.links[0]->length;  // 上臂长度
    double l2 = arm.links[1]->length;  // 下臂长度
    
    // 关节角度
    double theta1 = joint_angles[0];  // 基座关节 (绕Z轴)
    double theta2 = joint_angles[1];  // 肩部关节 (绕Y轴)
    double theta3 = joint_angles[2];  // 肘部关节 (绕Y轴)
    double theta4 = joint_angles[3];  // 腕部关节 (绕Y轴) - 简化模型中不影响位置
    
    // 计算末端位置
    Eigen::Vector3d end_pos;
    
    // 正运动学公式 (简化的3DOF模型)
    double x = (l1 * std::cos(theta2) + l2 * std::cos(theta2 + theta3)) * std::cos(theta1);
    double y = (l1 * std::cos(theta2) + l2 * std::cos(theta2 + theta3)) * std::sin(theta1);
    double z = l1 * std::sin(theta2) + l2 * std::sin(theta2 + theta3);
    
    end_pos << x, y, z;
    return end_pos;
}

bool RobotKinematics::arm_inverse_kinematics(const Eigen::Vector3d& target_pos, std::vector<double>& joint_angles)
{
    const auto& arm = robot_model_->get_arm();
    joint_angles.resize(4, 0.0);  // 4个关节
    
    // 连杆长度
    double l1 = arm.links[0]->length;
    double l2 = arm.links[1]->length;
    
    // 目标位置坐标
    double x = target_pos.x();
    double y = target_pos.y();
    double z = target_pos.z();
    
    // 计算手腕位置到基座的距离
    double r = std::sqrt(x*x + y*y);
    double d = std::sqrt(r*r + z*z);
    
    // 检查工作空间
    if (d > l1 + l2 || d < std::fabs(l1 - l2))
    {
        std::cerr << "Target position out of arm workspace" << std::endl;
        return false;
    }
    
    // 计算关节角度 (逆运动学解)
    // 基座关节角度
    joint_angles[0] = std::atan2(y, x);
    
    // 肩部和肘部关节角度 (使用余弦定理)
    double cos_theta3 = (d*d - l1*l1 - l2*l2) / (2*l1*l2);
    cos_theta3 = std::clamp(cos_theta3, -1.0, 1.0);  // 防止数值误差导致的问题
    joint_angles[2] = std::acos(cos_theta3);
    
    double alpha = std::atan2(z, r);
    double beta = std::asin((l2 * std::sin(joint_angles[2])) / d);
    joint_angles[1] = alpha + beta;
    
    // 腕部关节 (简化为0)
    joint_angles[3] = 0.0;
    
    // 检查关节角度限制
    for (size_t i = 0; i < joint_angles.size(); ++i)
    {
        if (joint_angles[i] < arm.joints[i]->min_angle || 
            joint_angles[i] > arm.joints[i]->max_angle)
        {
            std::cerr << "Joint " << i << " angle out of range" << std::endl;
            return false;
        }
    }
    
    return true;
}

Eigen::Vector3d RobotKinematics::leg_forward_kinematics(size_t leg_index, const std::vector<double>& joint_angles)
{
    const auto& leg = robot_model_->get_leg(leg_index);
    
    // 检查关节角度数量是否正确
    if (joint_angles.size() != leg.joints.size())
    {
        throw std::invalid_argument("Invalid number of joint angles for leg");
    }
    
    // 连杆长度
    double l1 = leg.links[0]->length;  // 大腿长度
    double l2 = leg.links[1]->length;  // 小腿长度
    
    // 关节角度
    double theta1 = joint_angles[0];  // 髋关节
    double theta2 = joint_angles[1];  // 膝关节
    
    // 根据腿部位置调整坐标系
    double sign = (leg.side.find("right") != std::string::npos) ? -1.0 : 1.0;
    
    // 计算足部位置
    Eigen::Vector3d foot_pos;
    
    // 正运动学计算
    foot_pos.x() = sign * (l1 * std::cos(theta1) + l2 * std::cos(theta1 + theta2));
    foot_pos.y() = 0.0;  // 简化模型
    foot_pos.z() = -(l1 * std::sin(theta1) + l2 * std::sin(theta1 + theta2));  // Z轴向下为正
    
    return foot_pos;
}

bool RobotKinematics::leg_inverse_kinematics(size_t leg_index, const Eigen::Vector3d& target_pos, std::vector<double>& joint_angles)
{
    const auto& leg = robot_model_->get_leg(leg_index);
    joint_angles.resize(2, 0.0);  // 每条腿2个关节
    
    // 连杆长度
    double l1 = leg.links[0]->length;  // 大腿长度
    double l2 = leg.links[1]->length;  // 小腿长度
    
    // 根据腿部位置调整坐标系
    double sign = (leg.side.find("right") != std::string::npos) ? -1.0 : 1.0;
    
    // 目标位置坐标 (调整符号)
    double x = sign * target_pos.x();
    double z = -target_pos.z();  // 转换为腿部坐标系
    
    // 计算髋关节到目标点的距离
    double d = std::sqrt(x*x + z*z);
    
    // 检查工作空间
    if (d > l1 + l2 || d < std::fabs(l1 - l2))
    {
        std::cerr << "Target position out of leg workspace" << std::endl;
        return false;
    }
    
    // 计算关节角度 (逆运动学解)
    double cos_theta2 = (d*d - l1*l1 - l2*l2) / (2*l1*l2);
    cos_theta2 = std::clamp(cos_theta2, -1.0, 1.0);  // 防止数值误差导致的问题
    joint_angles[1] = std::acos(cos_theta2);
    
    double alpha = std::atan2(z, x);
    double beta = std::asin((l2 * std::sin(joint_angles[1])) / d);
    joint_angles[0] = alpha - beta;
    
    // 检查关节角度限制
    for (size_t i = 0; i < joint_angles.size(); ++i)
    {
        if (joint_angles[i] < leg.joints[i]->min_angle || 
            joint_angles[i] > leg.joints[i]->max_angle)
        {
            std::cerr << "Leg joint " << i << " angle out of range: " 
                      << joint_angles[i] << " (min: " << leg.joints[i]->min_angle 
                      << ", max: " << leg.joints[i]->max_angle << ")" << std::endl;
            return false;
        }
    }
    
    return true;
}

}  // namespace robot_model
