/*
 * LQR控制器实现
 */

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

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

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

/* 从MATLAB仿真get_k.m得到的默认增益矩阵多项式拟合系数 */
/* 这些系数是通过对不同腿长的LQR增益进行3阶多项式拟合得到的 */
static LQR_GainMatrix_t default_gain_matrix = {
    /* K矩阵第一行 - 轮毂力矩T的增益系数 */
    .k11_coeff = {0.0f, -2845.3f, 1899.4f, -123.8f},      /* theta */
    .k12_coeff = {0.0f, -89.7f, 61.2f, -4.8f},            /* d_theta */
    .k13_coeff = {0.0f, 5479.2f, -3298.6f, 489.8f},       /* x */
    .k14_coeff = {0.0f, 312.4f, -178.9f, 34.2f},          /* d_x */
    .k15_coeff = {0.0f, -31250.0f, 18750.0f, -3125.0f},   /* phi */
    .k16_coeff = {0.0f, -89.7f, 61.2f, -4.8f},            /* d_phi */
    
    /* K矩阵第二行 - 髋关节力矩Tp的增益系数 */
    .k21_coeff = {0.0f, 486.1f, -324.1f, 21.6f},          /* theta */
    .k22_coeff = {0.0f, 15.3f, -10.4f, 0.8f},             /* d_theta */
    .k23_coeff = {0.0f, -935.4f, 562.2f, -83.5f},         /* x */
    .k24_coeff = {0.0f, -53.3f, 30.5f, -5.8f},            /* d_x */
    .k25_coeff = {0.0f, 5333.3f, -3200.0f, 533.3f},       /* phi */
    .k26_coeff = {0.0f, 15.3f, -10.4f, 0.8f},             /* d_phi */
};

/* Private function prototypes ---------------------------------------------- */
static int8_t LQR_CalculateErrorState(LQR_Controller_t *lqr);
static int8_t LQR_ApplyControlLimits(LQR_Controller_t *lqr);

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

int8_t LQR_Init(LQR_Controller_t *lqr, float max_wheel_torque, float max_joint_torque) {
    if (lqr == NULL) {
        return -1;
    }
    
    /* 清零结构体 */
    memset(lqr, 0, sizeof(LQR_Controller_t));
    
    /* 设置控制限制 */
    lqr->param.max_wheel_torque = max_wheel_torque;
    lqr->param.max_joint_torque = max_joint_torque;
    
    /* 设置默认系统物理参数（从MATLAB仿真get_k_length.m获取） */
    lqr->param.R = 0.072f;          /* 驱动轮半径 */
    lqr->param.l = 0.03f;           /* 机体重心到转轴距离 */
    lqr->param.mw = 0.60f;          /* 驱动轮质量 */
    lqr->param.mp = 1.8f;           /* 摆杆质量 */
    lqr->param.M = 1.8f;            /* 机体质量 */
    lqr->param.g = 9.8f;            /* 重力加速度 */
    
    /* 计算转动惯量 */
    lqr->param.Iw = lqr->param.mw * lqr->param.R * lqr->param.R;
    lqr->param.IM = lqr->param.M * (0.3f * 0.3f + 0.12f * 0.12f) / 12.0f;
    
    /* 设置默认增益矩阵 */
    lqr->gain_matrix = &default_gain_matrix;
    
    /* 设置默认目标状态（平衡状态） */
    memset(&lqr->param.target_state, 0, sizeof(LQR_State_t));
    
    /* 初始化当前腿长为中等值 */
    lqr->current_leg_length = 0.25f;
    
    /* 计算初始增益矩阵 */
    LQR_CalculateGainMatrix(lqr, lqr->current_leg_length);
    
    lqr->initialized = true;
    
    return 0;
}

int8_t LQR_SetGainMatrix(LQR_Controller_t *lqr, LQR_GainMatrix_t *gain_matrix) {
    if (lqr == NULL || gain_matrix == NULL) {
        return -1;
    }
    
    lqr->gain_matrix = gain_matrix;
    
    /* 重新计算增益矩阵 */
    return LQR_CalculateGainMatrix(lqr, lqr->current_leg_length);
}

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

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

int8_t LQR_CalculateGainMatrix(LQR_Controller_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->param.L = leg_length / 2.0f;           /* 摆杆重心到驱动轮轴距离 */
    lqr->param.LM = leg_length / 2.0f;          /* 摆杆重心到其转轴距离 */
    
    /* 计算摆杆转动惯量 */
    float leg_total_length = lqr->param.L + lqr->param.LM;
    lqr->param.Ip = lqr->param.mp * (leg_total_length * leg_total_length + 0.05f * 0.05f) / 12.0f;
    
    /* 使用多项式拟合计算当前增益矩阵 */
    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_Controller_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 LQR_ApplyControlLimits(lqr);
}

int8_t LQR_GetOutput(LQR_Controller_t *lqr, LQR_Input_t *output) {
    if (lqr == NULL || output == NULL) {
        return -1;
    }
    
    *output = lqr->control_output;
    return 0;
}

int8_t LQR_Reset(LQR_Controller_t *lqr) {
    if (lqr == NULL) {
        return -1;
    }
    
    /* 清零状态和输出 */
    memset(&lqr->current_state, 0, sizeof(LQR_State_t));
    memset(&lqr->error_state, 0, sizeof(LQR_State_t));
    memset(&lqr->control_output, 0, sizeof(LQR_Input_t));
    memset(&lqr->param.target_state, 0, sizeof(LQR_State_t));
    
    /* 重置限幅标志 */
    lqr->wheel_torque_limited = false;
    lqr->joint_torque_limited = false;
    
    return 0;
}

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];
}

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

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

static int8_t LQR_ApplyControlLimits(LQR_Controller_t *lqr) {
    if (lqr == NULL) {
        return -1;
    }
    
    /* 重置限幅标志 */
    lqr->wheel_torque_limited = false;
    lqr->joint_torque_limited = false;
    
    /* 限制轮毂力矩 */
    if (fabsf(lqr->control_output.T) > lqr->param.max_wheel_torque) {
        lqr->control_output.T = LQR_LIMIT(lqr->control_output.T, 
                                         -lqr->param.max_wheel_torque, 
                                          lqr->param.max_wheel_torque);
        lqr->wheel_torque_limited = true;
    }
    
    /* 限制髋关节力矩 */
    if (fabsf(lqr->control_output.Tp) > lqr->param.max_joint_torque) {
        lqr->control_output.Tp = LQR_LIMIT(lqr->control_output.Tp,
                                          -lqr->param.max_joint_torque,
                                           lqr->param.max_joint_torque);
        lqr->joint_torque_limited = true;
    }
    
    return 0;
}