/*
 * 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 ----------------------------------------------------------- */

typedef struct {
    float theta;      /* 摆杆角度 */
    float d_theta;    /* 摆杆角速度 */
    float x;          /* 驱动轮位移 */
    float d_x;        /* 驱动轮速度 */
    float phi;        /* 机体俯仰角 */
    float d_phi;      /* 机体俯仰角速度 */
} LQR_State_t;

typedef struct {
    float T;          /* 轮毂力矩 (N·m) */
    float Tp;         /* 髋关节力矩 (N·m) */
} LQR_Output_t;

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 {
    LQR_GainMatrix_t *gain_matrix; /* 增益矩阵参数指针 */
    
    LQR_State_t current_state;   /* 当前状态 */
    LQR_State_t target_state;    /* 目标状态 */
    LQR_State_t error_state;     /* 状态误差 */

    LQR_Output_t control_output;  /* 控制输出 */
    
    /* 运行时变量 */
    float current_leg_length;    /* 当前腿长 */
    float K_matrix[LQR_INPUT_DIM][LQR_STATE_DIM]; /* 当前增益矩阵 */
    
    bool initialized;            /* 初始化标志 */
    
} LQR_t;

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

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


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

/**
 * @brief 设置目标状态
 * 
 * @param lqr LQR控制器结构体指针
 * @param target_state 目标状态
 * @return int8_t 0-成功, 其他-失败
 */
int8_t LQR_SetTargetState(LQR_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_t *lqr, float leg_length);

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

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

#ifdef __cplusplus
}
#endif