#include "robot_model.hpp"
#include <iostream>
#include <cmath>

namespace robot_model
{

    RobotModel::RobotModel()
    {
        initialize_model();
    }

    void RobotModel::initialize_model()
    {
        // 初始化机械臂（4个关节，2个连杆）
        // 添加关节
        arm_.joints.push_back(std::make_shared<Joint>("arm_base_joint", -M_PI/2, M_PI/2, 10.0));
        arm_.joints.push_back(std::make_shared<Joint>("arm_shoulder_joint", -M_PI/2, M_PI/2, 8.0));
        arm_.joints.push_back(std::make_shared<Joint>("arm_elbow_joint", 0, M_PI/2, 6.0));
        arm_.joints.push_back(std::make_shared<Joint>("arm_wrist_joint", -M_PI/4, M_PI/4, 4.0));
        
        // 添加连杆
        arm_.links.push_back(std::make_shared<Link>("arm_upper_link", 0.3, 1.5));
        arm_.links.push_back(std::make_shared<Link>("arm_lower_link", 0.25, 1.0));

        // 初始化腿部（4条腿，每条腿2个连杆）
        std::vector<std::string> leg_sides = {"front_left", "front_right", "rear_left", "rear_right"};
        
        for (const auto & side : leg_sides)
        {
            Leg leg;
            leg.side = side;
            
            // 为每条腿添加关节
            leg.joints.push_back(std::make_shared<Joint>(side + "_hip_joint", -M_PI/4, M_PI/4, 15.0));
            leg.joints.push_back(std::make_shared<Joint>(side + "_knee_joint", 0, M_PI/2, 12.0));
            
            // 为每条腿添加连杆
            leg.links.push_back(std::make_shared<Link>(side + "_thigh_link", 0.2, 2.0));
            leg.links.push_back(std::make_shared<Link>(side + "_shin_link", 0.2, 1.5));
            
            legs_.push_back(leg);
        }
    }

    const Arm & RobotModel::get_arm(size_t index) const
    {
        if (index >= arms_.size())
        {
            throw std::out_of_range("Arm index out of range");
        }
        return arms_[index];
    }

    const Leg & RobotModel::get_leg(size_t index) const
    {
        if (index >= legs_.size())
        {
            throw std::out_of_range("Leg index out of range");
        }
        return legs_[index];
    }

    bool RobotModel::set_joint_angle(const std::string & joint_name, double angle)
    {
        // 检查机械臂关节
        for (auto & arm : arms_)
        {
            for (const auto & joint : arm.joints)
            {
                if (joint->name == joint_name)
                {
                    // 检查角度限制
                    if (angle >= joint->min_angle && angle <= joint->max_angle)
                    {
                        joint->angle = angle;
                        return true;
                    }
                    return false;
                }
            }
        }
        
        // 检查腿部关节
        for (auto & leg : legs_)
        {
            for (const auto & joint : leg.joints)
            {
                if (joint->name == joint_name)
                {
                    // 检查角度限制
                    if (angle >= joint->min_angle && angle <= joint->max_angle)
                    {
                        joint->angle = angle;
                        return true;
                    }
                    return false;
                }
            }
        }
        
        // 未找到关节
        return false;
    }

    double RobotModel::get_joint_angle(const std::string & joint_name) const
    {
        // 检查机械臂关节
        for (auto & arm : arms_)
        {
            for (const auto & joint : arm.joints)
            {
                if (joint->name == joint_name)
                {
                    return joint->angle;
                }
            }
        }
        
        // 检查腿部关节
        for (const auto & leg : legs_)
        {
            for (const auto & joint : leg.joints)
            {
                if (joint->name == joint_name)
                {
                    return joint->angle;
                }
            }
        }
        
        // 未找到关节
        return 0.0;
    }

    void RobotModel::print_model_info() const
    {
        std::cout << "=== Robot Model Information ===" << std::endl;
        
        // 打印机械臂信息
        std::cout << "\nArm:" << std::endl;
        std::cout << "  Joints:" << std::endl;
        for (size_t i = 0; i < arms_.size(); ++i)
        {
            const auto & arm = arms_[i];
            std::cout << "  Leg " << i << " (" << arm.side << "):" << std::endl;
            std::cout << "    Joints:" << std::endl;
            for (const auto & joint : arm.joints)
            {
                std::cout << "    " << joint->name << ": "
                        << "Angle=" << joint->angle << " rad, "
                        << "Range=[" << joint->min_angle << ", " << joint->max_angle << "] rad" << std::endl;
            }
            std::cout << "  Links:" << std::endl;
            for (const auto & link : arm.links)
            {
                std::cout << "    " << link->name << ": "
                        << "Length=" << link->length << " m, "
                        << "Mass=" << link->mass << " kg" << std::endl;
            }
        }
        
        // 打印腿部信息
        std::cout << "\nLegs:" << std::endl;
        for (size_t i = 0; i < legs_.size(); ++i)
        {
            const auto & leg = legs_[i];
            std::cout << "  Leg " << i << " (" << leg.side << "):" << std::endl;
            std::cout << "    Joints:" << std::endl;
            for (const auto & joint : leg.joints)
            {
                std::cout << "      " << joint->name << ": "
                        << "Angle=" << joint->angle << " rad, "
                        << "Range=[" << joint->min_angle << ", " << joint->max_angle << "] rad" << std::endl;
            }
            std::cout << "    Links:" << std::endl;
            for (const auto & link : leg.links)
            {
                std::cout << "      " << link->name << ": "
                        << "Length=" << link->length << " m, "
                        << "Mass=" << link->mass << " kg" << std::endl;
            }
        }
    }

}  // namespace robot_model
