#include "simple_quadruped_control/simple_control.hpp"

// tau = tff + kp(q - q0) + kd(dq - dq0)
// tau :输出力矩
// tff :前馈力矩
// q：期望角度位置
// q0：当前角度位置
// dq 期望角速度
// dq0 当前角速度
// kp：位置刚度
// kd：速度阻尼

QuadrupedSubscriber::QuadrupedSubscriber() : Node("simple_control")
{
    // 订阅 Ps2Data 话题
    ps2_data_subscriber_ = this->create_subscription<rc_msgs::msg::Ps2Data>(
        "ps2_data", 10,
        std::bind(&QuadrupedSubscriber::ps2_data_callback, this, std::placeholders::_1));

    // 发布 MotorCmds 话题
    motor_cmds_publisher_ = this->create_publisher<rc_msgs::msg::MotorCmds>(
        "/LF/motor_cmds", 10);

    // 订阅 DataMotors 话题
    data_motors_subscriber_ = this->create_subscription<rc_msgs::msg::DataMotors>(
        "/LF/data_motors", 10,
        std::bind(&QuadrupedSubscriber::data_motors_callback, this, std::placeholders::_1));

    // 初始化机械零点
    mech_zero.motors[0] = 0.7438;
    mech_zero.motors[1] = 0.5693;
    mech_zero.motors[2] = 0.6547;
}

void QuadrupedSubscriber::ps2_data_callback(const rc_msgs::msg::Ps2Data::SharedPtr msg)
{
    ps2.lx = msg->lx;
    ps2.ly = msg->ly;
    ps2.rx = msg->rx;
    ps2.ry = msg->ry;

    ps2.up_down = msg->up_down;
    ps2.left_right = msg->left_right;

    ps2.l1 = msg->l1;
    ps2.l2 = msg->l2;
    ps2.r1 = msg->r1;
    ps2.r2 = msg->r2;

    ps2.x = msg->x;
    ps2.y = msg->y;
    ps2.a = msg->a;
    ps2.b = msg->b;

    ps2.select = msg->select;
    ps2.start = msg->start;

    // 回零
    if (ps2.l1)
    {
        for (int i = 0; i < 3; ++i)
        {   // 重置commonds
            cmds.motors[i].pos = mech_zero.motors[i];
            cmds.motors[i].vw = 0.0;
            cmds.motors[i].tau = 0.0;
            cmds.motors[i].kp = 1.2;
            cmds.motors[i].kd = 0.05;
            // 限制力矩输出
            motor_cmds.motors[i].pos = cmds.motors[i].pos;
            motor_cmds.motors[i].vw = cmds.motors[i].vw;
            motor_cmds.motors[i].tau = cmds.motors[i].tau;
            motor_cmds.motors[i].kp = cmds.motors[i].kp;
            motor_cmds.motors[i].kd = cmds.motors[i].kd;
            motor_cmds.motors[i].tau = limit_torque(motor_cmds.motors[i], data_motors.motors[i], 0.5);
        }
        publish_motor_cmds();
        RCLCPP_INFO(this->get_logger(), "Zeroing motors");
        return;
    }
    if (ps2.l2)
    {
        
            cmds.motors[0].pos = cmds.motors[0].pos + ps2.lx * 0.01;
            cmds.motors[0].vw = 0.0;
            cmds.motors[0].tau = 0.0;
            cmds.motors[0].kp = 1.2;
            cmds.motors[0].kd = 0.05;

            cmds.motors[1].pos = cmds.motors[1].pos + ps2.ly * 0.01;
            cmds.motors[1].vw = 0.0;
            cmds.motors[1].tau = 0.0;
            cmds.motors[1].kp = 1.2;
            cmds.motors[1].kd = 0.05;

            cmds.motors[2].pos = cmds.motors[2].pos + ps2.ly * 0.01;
            cmds.motors[2].vw = 0.0;
            cmds.motors[2].tau = 0.0;
            cmds.motors[2].kp = 1.2;
            cmds.motors[2].kd = 0.05;
            // 限制力矩输出
        for (int i = 0; i < 3; ++i)
        {
            motor_cmds.motors[i].pos = cmds.motors[i].pos;
            motor_cmds.motors[i].vw = cmds.motors[i].vw;
            motor_cmds.motors[i].tau = cmds.motors[i].tau;
            motor_cmds.motors[i].kp = cmds.motors[i].kp;
            motor_cmds.motors[i].kd = cmds.motors[i].kd;
            motor_cmds.motors[i].tau = limit_torque(motor_cmds.motors[i], data_motors.motors[i], 0.5);
        }
        publish_motor_cmds();
        // RCLCPP_INFO(this->get_logger(), "Zeroing motors");
        return;
    }
}

void QuadrupedSubscriber::data_motors_callback(const rc_msgs::msg::DataMotors::SharedPtr msg)
{
    data_motors.motors[0].pos = msg->data_motor_0.pos;
    data_motors.motors[0].vw = msg->data_motor_0.vw;
    data_motors.motors[0].tau = msg->data_motor_0.tau;

    data_motors.motors[1].pos = msg->data_motor_1.pos;
    data_motors.motors[1].vw = msg->data_motor_1.vw;
    data_motors.motors[1].tau = msg->data_motor_1.tau;

    data_motors.motors[2].pos = msg->data_motor_2.pos;
    data_motors.motors[2].vw = msg->data_motor_2.vw;
    data_motors.motors[2].tau = msg->data_motor_2.tau;

    // RCLCPP_INFO(this->get_logger(), "Updated DataMotors for motors 0, 1, 2");
}

void QuadrupedSubscriber::publish_motor_cmds()
{
    auto motor_cmds_msg = rc_msgs::msg::MotorCmds();

    // // 设置每个电机的命令
    // for (int i = 0; i < 3; ++i)
    // {
    //     motor_cmds_msg.motor_cmds[i].tau = motor_cmds.motors[i].tau;
    //     motor_cmds_msg.motor_cmds[i].vw = motor_cmds.motors[i].vw;
    //     motor_cmds_msg.motor_cmds[i].pos = motor_cmds.motors[i].pos;
    //     motor_cmds_msg.motor_cmds[i].kp = motor_cmds.motors[i].kp;
    //     motor_cmds_msg.motor_cmds[i].kd = motor_cmds.motors[i].kd;
    // }
        // 设置每个电机的命令
        motor_cmds_msg.motor_cmd_0.tau = motor_cmds.motors[0].tau;
        motor_cmds_msg.motor_cmd_0.vw = motor_cmds.motors[0].vw;
        motor_cmds_msg.motor_cmd_0.pos = motor_cmds.motors[0].pos;
        motor_cmds_msg.motor_cmd_0.kp = motor_cmds.motors[0].kp;
        motor_cmds_msg.motor_cmd_0.kd = motor_cmds.motors[0].kd;
    
        motor_cmds_msg.motor_cmd_1.tau = motor_cmds.motors[1].tau;
        motor_cmds_msg.motor_cmd_1.vw = motor_cmds.motors[1].vw;
        motor_cmds_msg.motor_cmd_1.pos = motor_cmds.motors[1].pos;
        motor_cmds_msg.motor_cmd_1.kp = motor_cmds.motors[1].kp;
        motor_cmds_msg.motor_cmd_1.kd = motor_cmds.motors[1].kd;
    
        motor_cmds_msg.motor_cmd_2.tau = motor_cmds.motors[2].tau;
        motor_cmds_msg.motor_cmd_2.vw = motor_cmds.motors[2].vw;
        motor_cmds_msg.motor_cmd_2.pos = motor_cmds.motors[2].pos;
        motor_cmds_msg.motor_cmd_2.kp = motor_cmds.motors[2].kp;
        motor_cmds_msg.motor_cmd_2.kd = motor_cmds.motors[2].kd;
    

    // 发布 MotorCmds 消息
    motor_cmds_publisher_->publish(motor_cmds_msg);
}

// 限制输出力矩的函数
float QuadrupedSubscriber::limit_torque(const motor_cmd_t &motor_cmd, const data_motor_t &data_motor, float limit)
{
    // 读取结构体中的参数
    float tff = motor_cmd.tau;  // 前馈力矩
    float kp = motor_cmd.kp;    // 位置刚度
    float kd = motor_cmd.kd;    // 速度阻尼
    float q = motor_cmd.pos*6.33;    // 期望角度位置
    float q0 = data_motor.pos*6.33;  // 当前角度位置
    float dq = motor_cmd.vw*6.33;    // 期望角速度
    float dq0 = data_motor.vw*6.33;  // 当前角速度

    // 计算输出力矩
    float tau = tff + kp * (q - q0) + kd * (dq - dq0);

    // 如果力矩超过限制值，调整前馈力矩
    if (tau > limit)
    {
        tff -= (tau - limit) / 2;
        tau = limit;
        tff = tau - kp * (q - q0) - kd * (dq - dq0);
        return tff;
    }else if (tau < -limit)
    {
        tff += (tau + limit) / 2;
        tau = -limit;
        tff = tau - kp * (q - q0) - kd * (dq - dq0);
        return tff;
    }
    return tff;
}

int main(int argc, char *argv[])
{
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<QuadrupedSubscriber>());
    rclcpp::shutdown();
    return 0;
}