#include "balance.h"
#include "stdint.h"
#include "arm_math.h"

/**
 * @brief 根据状态反馈计算当前腿长,查表获得LQR的反馈增益,并列式计算LQR的输出
 * @note 得到的腿部力矩输出还要经过综合运动控制系统补偿后映射为两个关节电机输出
 *
 */
static void CalcLQR(LinkNPodParam *p, ChassisParam *chassis)
{
    static float k[12][3] = {
        {91.443258f, -113.149301f, -8.780431f},
        {2.186397f, -6.908826f, -0.171192f},
        {47.943203f, -39.423913f, -13.297140f},
        {30.671413f, -28.308799f, -9.363969f},
        {231.778706f, -224.742110f, 68.974495f},
        {18.910995f, -20.079072f, 7.490403f},
        {50.114072f, -58.670210f, 23.856863f},
        {3.273746f, -3.321150f, 1.267603f},
        {140.969395f, -137.481822f, 42.507159f},
        {93.859327f, -90.754295f, 27.811839f},
        {-283.581279f, 232.017305f, 87.865897f},
        {-28.193619f, 23.802666f, 3.963794f},
    };
    float T[2] = {0};       // 0 T_wheel  1 T_hip
    float l = p->leg_len;
    float lsqr = l * l;

    uint8_t i, j;

    // 离地时轮子输出置0
    i = 0; j = i * 6;
    T[i] = p->fly_flag ? 0 :
           ((k[j + 0][0] * lsqr + k[j + 0][1] * l + k[j + 0][2]) * -p->theta +
            (k[j + 1][0] * lsqr + k[j + 1][1] * l + k[j + 1][2]) * -p->theta_w + 
            (k[j + 2][0] * lsqr + k[j + 2][1] * l + k[j + 2][2]) * (chassis->target_dist - chassis->dist) +
            (k[j + 3][0] * lsqr + k[j + 3][1] * l + k[j + 3][2]) * (chassis->target_v - chassis->vel) +
            (k[j + 4][0] * lsqr + k[j + 4][1] * l + k[j + 4][2]) * -chassis->pitch +
            (k[j + 5][0] * lsqr + k[j + 5][1] * l + k[j + 5][2]) * -chassis->pitch_w);
    
    // 离地时关节输出仅保留 theta 和 theta_dot，保证滞空时腿部竖直
    i = 1; j = i * 6;
    T[i] = (k[j + 0][0] * lsqr + k[j + 0][1] * l + k[j + 0][2]) * -p->theta +
           (k[j + 1][0] * lsqr + k[j + 1][1] * l + k[j + 1][2]) * -p->theta_w + (p->fly_flag ? 0 :
           ((k[j + 2][0] * lsqr + k[j + 2][1] * l + k[j + 2][2]) * (chassis->target_dist - chassis->dist) +
            (k[j + 3][0] * lsqr + k[j + 3][1] * l + k[j + 3][2]) * (chassis->target_v - chassis->vel) +
            (k[j + 4][0] * lsqr + k[j + 4][1] * l + k[j + 4][2]) * -chassis->pitch +
            (k[j + 5][0] * lsqr + k[j + 5][1] * l + k[j + 5][2]) * -chassis->pitch_w));
    
    p->T_wheel = T[0];
    p->T_hip = T[1];
}