#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] = {62.680622,-74.772126,-13.135672,
                            1.620796,-4.331826,-0.454705,
                            32.093563,-25.558681,-16.605856,
                            19.128242,-17.562747,-10.815380,
                            225.373594,-201.324771,53.236052,
                            13.575849,-13.013546,3.845604,
                            76.704297,-72.201666,20.359891,
                            4.520141,-4.051799,1.189345,
                            139.418565,-125.750074,34.110069,
                            88.295852,-79.170124,21.430975,
                            -163.725633,127.852860,110.004619,
                            -9.557678,7.476790,4.655220};
    float T[2] = {0};       // 0 T_wheel  1 T_hip
    float l = p->leg_len;
    float lsqr = l * l;

    for (uint8_t i = 0; i < 2; ++i)
    {
        uint8_t j = i * 6;

        if(i == 0)      // 离地时仅对速度闭环，保证落地时轮速与机体速度一致
        {
            T[i] = (k[j + 3][0] * lsqr + k[j + 3][1] * l + k[j + 3][2]) * (chassis->target_v - chassis->vel) + (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 + 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 ));
        }
        else if(i == 1)     // 离地时关节输出仅保留 theta 和 theta_dot，保证滞空时腿部竖直
        {
            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];
}