/*
 * LQR控制器实现：单腿建模
 */

#include "lqr.h"
#include <string.h>
#include <stddef.h>

/* Private macros ----------------------------------------------------------- */
#define LQR_LIMIT(x, min, max) ((x) < (min) ? (min) : ((x) > (max) ? (max) : (x)))

/* Private variables -------------------------------------------------------- */

/* Private function prototypes ---------------------------------------------- */
static int8_t LQR_CalculateErrorState(LQR_t *lqr);
static float LQR_PolynomialCalc(const float *coeff, float leg_length);

/* Public functions --------------------------------------------------------- */

int8_t LQR_Init(LQR_t *lqr, LQR_GainMatrix_t *gain_matrix) {
    if (lqr == NULL || gain_matrix == NULL) {
        return -1;
    }
    
    /* 清零结构体 */
    memset(lqr, 0, sizeof(LQR_t));
    
    /* 设置增益矩阵 */
    lqr->gain_matrix = gain_matrix;
    
    /* 初始化当前腿长为中等值 */
    lqr->current_leg_length = 0.25f;
    
    /* 计算初始增益矩阵 */
    LQR_CalculateGainMatrix(lqr, lqr->current_leg_length);
    
    lqr->initialized = true;
    
    return 0;
}

int8_t LQR_UpdateState(LQR_t *lqr, const LQR_State_t *current_state) {
    if (lqr == NULL || current_state == NULL) {
        return -1;
    }
    
    /* 更新当前状态 */
    lqr->current_state = *current_state;
    
    /* 计算状态误差 */
    return LQR_CalculateErrorState(lqr);
}

int8_t LQR_SetTargetState(LQR_t *lqr, const LQR_State_t *target_state) {
    if (lqr == NULL || target_state == NULL) {
        return -1;
    }
    
    lqr->target_state = *target_state;
    
    /* 重新计算状态误差 */
    return LQR_CalculateErrorState(lqr);
}

int8_t LQR_CalculateGainMatrix(LQR_t *lqr, float leg_length) {
    if (lqr == NULL || lqr->gain_matrix == NULL) {
        return -1;
    }
    
    /* 限制腿长范围 */
    leg_length = LQR_LIMIT(leg_length, 0.1f, 0.4f);
    lqr->current_leg_length = leg_length;
    
    /* 使用多项式拟合计算当前增益矩阵 */
    lqr->K_matrix[0][0] = LQR_PolynomialCalc(lqr->gain_matrix->k11_coeff, leg_length);  /* K11: theta */
    lqr->K_matrix[0][1] = LQR_PolynomialCalc(lqr->gain_matrix->k12_coeff, leg_length);  /* K12: d_theta */
    lqr->K_matrix[0][2] = LQR_PolynomialCalc(lqr->gain_matrix->k13_coeff, leg_length);  /* K13: x */
    lqr->K_matrix[0][3] = LQR_PolynomialCalc(lqr->gain_matrix->k14_coeff, leg_length);  /* K14: d_x */
    lqr->K_matrix[0][4] = LQR_PolynomialCalc(lqr->gain_matrix->k15_coeff, leg_length);  /* K15: phi */
    lqr->K_matrix[0][5] = LQR_PolynomialCalc(lqr->gain_matrix->k16_coeff, leg_length);  /* K16: d_phi */
    
    lqr->K_matrix[1][0] = LQR_PolynomialCalc(lqr->gain_matrix->k21_coeff, leg_length);  /* K21: theta */
    lqr->K_matrix[1][1] = LQR_PolynomialCalc(lqr->gain_matrix->k22_coeff, leg_length);  /* K22: d_theta */
    lqr->K_matrix[1][2] = LQR_PolynomialCalc(lqr->gain_matrix->k23_coeff, leg_length);  /* K23: x */
    lqr->K_matrix[1][3] = LQR_PolynomialCalc(lqr->gain_matrix->k24_coeff, leg_length);  /* K24: d_x */
    lqr->K_matrix[1][4] = LQR_PolynomialCalc(lqr->gain_matrix->k25_coeff, leg_length);  /* K25: phi */
    lqr->K_matrix[1][5] = LQR_PolynomialCalc(lqr->gain_matrix->k26_coeff, leg_length);  /* K26: d_phi */
    
    return 0;
}

int8_t LQR_Control(LQR_t *lqr) {
    if (lqr == NULL || !lqr->initialized) {
        return -1;
    }
    
    /* LQR控制律: u = -K * x_error
     * 其中 u = [T; Tp], x_error = x_current - x_target
     */
    
    /* 计算轮毂力矩T */
    lqr->control_output.T = -(
        lqr->K_matrix[0][0] * lqr->error_state.theta +
        lqr->K_matrix[0][1] * lqr->error_state.d_theta +
        lqr->K_matrix[0][2] * lqr->error_state.x +
        lqr->K_matrix[0][3] * lqr->error_state.d_x +
        lqr->K_matrix[0][4] * lqr->error_state.phi +
        lqr->K_matrix[0][5] * lqr->error_state.d_phi
    );
    
    /* 计算髋关节力矩Tp */
    lqr->control_output.Tp = -(
        lqr->K_matrix[1][0] * lqr->error_state.theta +
        lqr->K_matrix[1][1] * lqr->error_state.d_theta +
        lqr->K_matrix[1][2] * lqr->error_state.x +
        lqr->K_matrix[1][3] * lqr->error_state.d_x +
        lqr->K_matrix[1][4] * lqr->error_state.phi +
        lqr->K_matrix[1][5] * lqr->error_state.d_phi
    );
    
    return 0;
}

int8_t LQR_Reset(LQR_t *lqr) {
    if (lqr == NULL) {
        return -1;
    }
    
    /* 清零状态和输出 */
    memset(&lqr->current_state, 0, sizeof(LQR_State_t));
    memset(&lqr->target_state, 0, sizeof(LQR_State_t));
    memset(&lqr->error_state, 0, sizeof(LQR_State_t));
    memset(&lqr->control_output, 0, sizeof(LQR_Output_t));
    
    return 0;
}

/* Private functions -------------------------------------------------------- */

static int8_t LQR_CalculateErrorState(LQR_t *lqr) {
    if (lqr == NULL) {
        return -1;
    }
    
    /* 计算状态误差 */
    lqr->error_state.theta = lqr->current_state.theta - lqr->target_state.theta;
    lqr->error_state.d_theta = lqr->current_state.d_theta - lqr->target_state.d_theta;
    lqr->error_state.x = lqr->current_state.x - lqr->target_state.x;
    lqr->error_state.d_x = lqr->current_state.d_x - lqr->target_state.d_x;
    lqr->error_state.phi = lqr->current_state.phi - lqr->target_state.phi;
    lqr->error_state.d_phi = lqr->current_state.d_phi - lqr->target_state.d_phi;
    
    return 0;
}

static float LQR_PolynomialCalc(const float *coeff, float leg_length) {
    if (coeff == NULL) {
        return 0.0f;
    }
    
    /* 计算3阶多项式: coeff[0]*L^3 + coeff[1]*L^2 + coeff[2]*L + coeff[3] */
    float L = leg_length;
    float L2 = L * L;
    float L3 = L2 * L;
    
    return coeff[0] * L3 + coeff[1] * L2 + coeff[2] * L + coeff[3];
}