/*
 * LQR控制器
 * 用于轮腿平衡机器人的线性二次调节器控制
 */

#pragma once

#ifdef __cplusplus
extern "C" {
#endif

/* Includes ----------------------------------------------------------------- */
#include <stdint.h>
#include <stdbool.h>
#include <math.h>

/* Exported constants ------------------------------------------------------- */
#define LQR_STATE_DIM  6    /* 状态向量维度: theta, d_theta, x, d_x, phi, d_phi */
#define LQR_INPUT_DIM  2    /* 输入向量维度: T(轮毂力矩), Tp(髋关节力矩) */
#define LQR_POLY_ORDER 4    /* 多项式拟合阶数 */

/* Exported types ----------------------------------------------------------- */

/**
 * @brief LQR状态向量
 * 状态定义：
 * theta:   摆杆与竖直方向夹角 (rad)
 * d_theta: 摆杆角速度 (rad/s)
 * x:       驱动轮位移 (m)
 * d_x:     驱动轮速度 (m/s)
 * phi:     机体与水平夹角 (rad)
 * d_phi:   机体角速度 (rad/s)
 */
typedef struct {
    float theta;      /* 摆杆角度 */
    float d_theta;    /* 摆杆角速度 */
    float x;          /* 驱动轮位移 */
    float d_x;        /* 驱动轮速度 */
    float phi;        /* 机体俯仰角 */
    float d_phi;      /* 机体俯仰角速度 */
} LQR_State_t;

/**
 * @brief LQR控制输入向量
 */
typedef struct {
    float T;          /* 轮毂力矩 (N·m) */
    float Tp;         /* 髋关节力矩 (N·m) */
} LQR_Input_t;

/**
 * @brief LQR增益矩阵参数
 * K矩阵的每个元素的多项式拟合系数
 * K(leg_length) = a[0]*L^3 + a[1]*L^2 + a[2]*L + a[3]
 */
typedef struct {
    /* K矩阵第一行(轮毂力矩T对应的增益) */
    float k11_coeff[LQR_POLY_ORDER];  /* K(1,1): theta */
    float k12_coeff[LQR_POLY_ORDER];  /* K(1,2): d_theta */
    float k13_coeff[LQR_POLY_ORDER];  /* K(1,3): x */
    float k14_coeff[LQR_POLY_ORDER];  /* K(1,4): d_x */
    float k15_coeff[LQR_POLY_ORDER];  /* K(1,5): phi */
    float k16_coeff[LQR_POLY_ORDER];  /* K(1,6): d_phi */
    
    /* K矩阵第二行(髋关节力矩Tp对应的增益) */
    float k21_coeff[LQR_POLY_ORDER];  /* K(2,1): theta */
    float k22_coeff[LQR_POLY_ORDER];  /* K(2,2): d_theta */
    float k23_coeff[LQR_POLY_ORDER];  /* K(2,3): x */
    float k24_coeff[LQR_POLY_ORDER];  /* K(2,4): d_x */
    float k25_coeff[LQR_POLY_ORDER];  /* K(2,5): phi */
    float k26_coeff[LQR_POLY_ORDER];  /* K(2,6): d_phi */
} LQR_GainMatrix_t;

/**
 * @brief LQR控制器参数
 */
typedef struct {
    /* 系统物理参数 */
    float R;          /* 驱动轮半径 (m) */
    float L;          /* 摆杆重心到驱动轮轴距离 (m) */
    float LM;         /* 摆杆重心到其转轴距离 (m) */
    float l;          /* 机体重心到其转轴距离 (m) */
    float mw;         /* 驱动轮质量 (kg) */
    float mp;         /* 摆杆质量 (kg) */
    float M;          /* 机体质量 (kg) */
    float Iw;         /* 驱动轮转动惯量 (kg·m²) */
    float Ip;         /* 摆杆绕质心转动惯量 (kg·m²) */
    float IM;         /* 机体绕质心转动惯量 (kg·m²) */
    float g;          /* 重力加速度 (m/s²) */
    
    /* 控制限制 */
    float max_wheel_torque;  /* 轮毂电机最大力矩 (N·m) */
    float max_joint_torque;  /* 关节电机最大力矩 (N·m) */
    
    /* 目标状态 */
    LQR_State_t target_state; /* 目标状态 */
} LQR_Param_t;

/**
 * @brief LQR控制器主结构体
 */
typedef struct {
    LQR_Param_t param;           /* 控制器参数 */
    LQR_GainMatrix_t *gain_matrix; /* 增益矩阵参数指针 */
    
    LQR_State_t current_state;   /* 当前状态 */
    LQR_State_t error_state;     /* 状态误差 */
    LQR_Input_t control_output;  /* 控制输出 */
    
    /* 运行时变量 */
    float current_leg_length;    /* 当前腿长 */
    float K_matrix[LQR_INPUT_DIM][LQR_STATE_DIM]; /* 当前增益矩阵 */
    
    bool initialized;            /* 初始化标志 */
    
    /* 限幅标志 */
    bool wheel_torque_limited;   /* 轮毂力矩是否被限幅 */
    bool joint_torque_limited;   /* 关节力矩是否被限幅 */
} LQR_Controller_t;

/* Exported functions prototypes -------------------------------------------- */

/**
 * @brief 初始化LQR控制器
 * 
 * @param lqr LQR控制器结构体指针
 * @param max_wheel_torque 轮毂电机最大力矩
 * @param max_joint_torque 关节电机最大力矩
 * @return int8_t 0-成功, 其他-失败
 */
int8_t LQR_Init(LQR_Controller_t *lqr, float max_wheel_torque, float max_joint_torque);

/**
 * @brief 设置LQR增益矩阵参数
 * 
 * @param lqr LQR控制器结构体指针  
 * @param gain_matrix 增益矩阵参数指针
 * @return int8_t 0-成功, 其他-失败
 */
int8_t LQR_SetGainMatrix(LQR_Controller_t *lqr, LQR_GainMatrix_t *gain_matrix);

/**
 * @brief 更新当前状态
 * 
 * @param lqr LQR控制器结构体指针
 * @param state 当前状态
 * @return int8_t 0-成功, 其他-失败
 */
int8_t LQR_UpdateState(LQR_Controller_t *lqr, const LQR_State_t *state);

/**
 * @brief 设置目标状态
 * 
 * @param lqr LQR控制器结构体指针
 * @param target_state 目标状态
 * @return int8_t 0-成功, 其他-失败
 */
int8_t LQR_SetTargetState(LQR_Controller_t *lqr, const LQR_State_t *target_state);

/**
 * @brief 根据当前腿长计算增益矩阵
 * 
 * @param lqr LQR控制器结构体指针
 * @param leg_length 当前腿长 (m)
 * @return int8_t 0-成功, 其他-失败
 */
int8_t LQR_CalculateGainMatrix(LQR_Controller_t *lqr, float leg_length);

/**
 * @brief 执行LQR控制计算
 * 
 * @param lqr LQR控制器结构体指针
 * @return int8_t 0-成功, 其他-失败
 */
int8_t LQR_Control(LQR_Controller_t *lqr);

/**
 * @brief 获取控制输出
 * 
 * @param lqr LQR控制器结构体指针
 * @param output 输出控制指令指针
 * @return int8_t 0-成功, 其他-失败
 */
int8_t LQR_GetOutput(LQR_Controller_t *lqr, LQR_Input_t *output);

/**
 * @brief 重置LQR控制器
 * 
 * @param lqr LQR控制器结构体指针
 * @return int8_t 0-成功, 其他-失败
 */
int8_t LQR_Reset(LQR_Controller_t *lqr);

/**
 * @brief 根据多项式系数计算增益值
 * 
 * @param coeff 多项式系数数组
 * @param leg_length 腿长
 * @return float 增益值
 */
float LQR_PolynomialCalc(const float *coeff, float leg_length);

#ifdef __cplusplus
}
#endif