#include "IMU.h"
#include <math.h>

//==================================================================================================
// 常量
//==================================================================================================

#ifndef PI
#define PI 3.141592654f
#endif

#define TWO_KP_DEF  (2.0f * 0.5f)   // Mahony 比例增益（默认 1.0）
#define TWO_KI_DEF  (2.0f * 0.0f)   // Mahony 积分增益（默认 0, 关闭）

//==================================================================================================
// Mahony 滤波器状态（静态，本文件内可见）
//==================================================================================================

static float inv_sample_freq;                           // 1 / 采样频率
static float two_kp = TWO_KP_DEF;                       // 比例增益
static float two_ki = TWO_KI_DEF;                       // 积分增益
static float q0 = 1.0f, q1 = 0.0f, q2 = 0.0f, q3 = 0.0f; // 姿态四元数
static float integral_fbx = 0.0f, integral_fby = 0.0f, integral_fbz = 0.0f; // 积分误差项

//==================================================================================================
// 陀螺仪零偏（由 imu_calibrate_gyro 填入）
//==================================================================================================

static float gyro_bias_x = 0.0f;
static float gyro_bias_y = 0.0f;
static float gyro_bias_z = 0.0f;

//==================================================================================================
// 姿态输出
//==================================================================================================

volatile imu_attitude_t imu_angle;

volatile int yaw_round_count = 0;

volatile static float yaw_angle_last = 0.0f;


//==================================================================================================
// 双二阶滤波器状态
//==================================================================================================

static biquad_state_t bq_gyro_x, bq_gyro_y, bq_gyro_z;
static biquad_state_t bq_accel_x, bq_accel_y, bq_accel_z;

//==================================================================================================
// 公开接口实现
//==================================================================================================

/** @brief 初始化 IMU
 *  @param sample_freq_hz 采样频率（Hz），等于定时器中断频率
 *  1. 初始化 ICM-42688 硬件
 *  2. 配置 Mahony 采样周期
 *  3. 初始化各轴陀螺仪/加速度计双二阶滤波器
 *  滤波器参数:
 *   陀螺仪 X/Y/Z: 低通 30/20/30Hz  → 滤除高频振动
 *   加速度计 X/Y/Z: 低通 10Hz      → 重力矢量提取（响应慢但稳）
 */
void imu_init(int sample_freq_hz)
{
    icm42688_init();

    inv_sample_freq = 1.0f / sample_freq_hz;

    biquad_init(&bq_gyro_x,  BIQUAD_LOWPASS, sample_freq_hz, 30, 0.7071f);
    biquad_init(&bq_gyro_y,  BIQUAD_LOWPASS, sample_freq_hz, 20, 0.5773f);
    biquad_init(&bq_gyro_z,  BIQUAD_LOWPASS, sample_freq_hz, 30, 0.5773f);
    biquad_init(&bq_accel_x, BIQUAD_LOWPASS, sample_freq_hz, 10, 0.5773f);
    biquad_init(&bq_accel_y, BIQUAD_LOWPASS, sample_freq_hz, 10, 0.5773f);
    biquad_init(&bq_accel_z, BIQUAD_LOWPASS, sample_freq_hz, 10, 0.5773f);
}

/** @brief 更新姿态（定时器 ISR 中按固定频率调用）
 *  数据流：传感器 → 零偏修正 → 双二阶滤波 → Mahony → 欧拉角 → Yaw 连续角度
 *  传感器坐标系直通，无坐标轴转换
 *  陀螺仪数据不进 Mahony 前只做偏置修正（Mahony 内部自带角速度积分处理）
 *  加速度数据经低通滤波后送入 Mahony 做重力矢量参考
 */
void imu_update(void)
{
    icm42688_data_t raw;
    icm42688_read_all(&raw);

    // 零偏修正
    raw.gyro_x -= gyro_bias_x;
    raw.gyro_y -= gyro_bias_y;
    raw.gyro_z -= gyro_bias_z;

    // 滤波（传感器坐标系，不做轴转换）
    imu_angle.gyro_x = biquad_apply(&bq_gyro_x, raw.gyro_x);
    imu_angle.gyro_y = biquad_apply(&bq_gyro_y, raw.gyro_y);
    imu_angle.gyro_z = biquad_apply(&bq_gyro_z, raw.gyro_z);

    float ax_f = biquad_apply(&bq_accel_x, raw.accel_x);
    float ay_f = biquad_apply(&bq_accel_y, raw.accel_y);
    float az_f = biquad_apply(&bq_accel_z, raw.accel_z);

    // Mahony AHRS 姿态解算
    mahony_ahrs_update(raw.gyro_x, raw.gyro_y, raw.gyro_z, ax_f, ay_f, az_f);

    // Yaw 连续角度跟踪（处理 ±180° 跳变）
    if (imu_angle.yaw - yaw_angle_last > 180.0f) {
        yaw_round_count--;
    } else if (imu_angle.yaw - yaw_angle_last < -180.0f) {
        yaw_round_count++;
    }
    yaw_angle_last = imu_angle.yaw;
}

/** @brief 陀螺仪零偏校准
 *  必须在设备完全静止时调用，采集 3 秒陀螺仪原始数据
 *  每秒采样 1000 点取均值，3 秒的均值再取平均，乘以比例系数得到 °/s 零偏
 *  校准结果自动存入 gyro_bias_x/y/z，在后续 imu_update() 中生效
 */
void imu_calibrate_gyro(void)
{
    #define CALIB_SECONDS 3
    #define CALIB_SAMPLES 1000

    static float gyro_history[3][CALIB_SECONDS];
    int32_t gyro_sum[3] = {0, 0, 0};
    uint32_t sec_cnt = 0, sample_cnt = 0;

    system_delay_ms(1000);
    system_delay_ms(1000);

    do {
        system_delay_us(1000);
        int16_t gx_r, gy_r, gz_r;
        icm42688_read_gyro_raw(&gx_r, &gy_r, &gz_r);

        if (sample_cnt < CALIB_SAMPLES) {
            gyro_sum[0] += gx_r;
            gyro_sum[1] += gy_r;
            gyro_sum[2] += gz_r;
            sample_cnt++;
        } else {
            gyro_history[0][sec_cnt] = (float)gyro_sum[0] / CALIB_SAMPLES;
            gyro_history[1][sec_cnt] = (float)gyro_sum[1] / CALIB_SAMPLES;
            gyro_history[2][sec_cnt] = (float)gyro_sum[2] / CALIB_SAMPLES;
            gyro_sum[0] = 0;
            gyro_sum[1] = 0;
            gyro_sum[2] = 0;
            sample_cnt = 0;
            sec_cnt++;
        }
    } while (sec_cnt < CALIB_SECONDS);

    float sum[3] = {0, 0, 0};
    float k = icm42688_gyro_scale() / CALIB_SECONDS;
    for (uint8_t i = 0; i < CALIB_SECONDS; i++) {
        sum[0] += gyro_history[0][i];
        sum[1] += gyro_history[1][i];
        sum[2] += gyro_history[2][i];
    }

    gyro_bias_x = k * sum[0];
    gyro_bias_y = k * sum[1];
    gyro_bias_z = k * sum[2];

}

//==================================================================================================
// 算法实现
//==================================================================================================

/** @brief Mahony AHRS 6轴姿态解算
 *  @param gx/gy/gz 陀螺仪角速度 (°/s)
 *  @param ax/ay/az 加速度计测量值 (g)
 *
 *  算法流程:
 *   1. 角速度转弧度，加速度归一化
 *   2. 由当前四元数估计重力方向
 *   3. 估计重力与测量重力的叉积 → PI 反馈修正陀螺仪
 *   4. 一阶龙格-库塔积分更新四元数
 *   5. 四元数归一化
 *   6. 提取欧拉角写入 imu_angle
 *
 *  参考: S.O.H. Madgwick "An efficient orientation filter for IMU and MARG sensor arrays"
 */
void mahony_ahrs_update(float gx, float gy, float gz, float ax, float ay, float az)
{
    float recip_norm;
    float halfvx, halfvy, halfvz;
    float halfex, halfey, halfez;
    float qa, qb, qc;

    // 仅当加速度计有读数时执行修正（静止或匀加速时加速度计才可靠）
    if (!((ax == 0.0f) && (ay == 0.0f) && (az == 0.0f))) {
        // 角速度 → rad/s
        gx *= ANG2RAD;
        gy *= ANG2RAD;
        gz *= ANG2RAD;

        // 加速度归一化
        recip_norm = inv_sqrt(ax * ax + ay * ay + az * az);
        ax *= recip_norm;
        ay *= recip_norm;
        az *= recip_norm;

        // 由当前四元数估计重力方向（世界系 Z 轴在传感器系的投影）
        halfvx = q1 * q3 - q0 * q2;
        halfvy = q0 * q1 + q2 * q3;
        halfvz = q0 * q0 - 0.5f + q3 * q3;

        // 叉积 = 估计方向 × 测量方向 → 姿态误差
        halfex = (ay * halfvz - az * halfvy);
        halfey = (az * halfvx - ax * halfvz);
        halfez = (ax * halfvy - ay * halfvx);

        // 积分反馈（默认 ki=0, 关闭）
        if (two_ki > 0.0f) {
            integral_fbx += two_ki * halfex * inv_sample_freq;
            integral_fby += two_ki * halfey * inv_sample_freq;
            integral_fbz += two_ki * halfez * inv_sample_freq;
            gx += integral_fbx;
            gy += integral_fby;
            gz += integral_fbz;
        } else {
            integral_fbx = 0.0f;
            integral_fby = 0.0f;
            integral_fbz = 0.0f;
        }

        // 比例反馈
        gx += two_kp * halfex;
        gy += two_kp * halfey;
        gz += two_kp * halfez;
    }

    // 四元数一阶积分 (RK1)
    gx *= (0.5f * inv_sample_freq);
    gy *= (0.5f * inv_sample_freq);
    gz *= (0.5f * inv_sample_freq);
    qa = q0; qb = q1; qc = q2;
    q0 += (-qb * gx - qc * gy - q3 * gz);
    q1 += ( qa * gx + qc * gz - q3 * gy);
    q2 += ( qa * gy - qb * gz + q3 * gx);
    q3 += ( qa * gz + qb * gy - qc * gx);

    // 四元数归一化
    recip_norm = inv_sqrt(q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3);
    q0 *= recip_norm;
    q1 *= recip_norm;
    q2 *= recip_norm;
    q3 *= recip_norm;

    // 旋转矩阵 → 欧拉角（传感器坐标系）
    float r11 = 2.0f * (q0 * q1 + q2 * q3);
    float r12 = 1.0f - 2.0f * (q1 * q1 + q2 * q2);
    float r21 = 2.0f * (q0 * q2 - q3 * q1);
    float r31 = 2.0f * (q0 * q3 + q1 * q2);
    float r32 = 1.0f - 2.0f * (q2 * q2 + q3 * q3);

    imu_angle.pitch =  RAD2ANG * asinf(r21);
    imu_angle.roll  =  RAD2ANG * atan2f(r11, r12);
    imu_angle.yaw   =  RAD2ANG * atan2f(r31, r32);
}

/** @brief 双二阶滤波器系数初始化
 *  @param state 滤波器状态
 *  @param type  滤波器类型
 *  @param fs    采样率 (Hz)
 *  @param fc    截止/中心频率 (Hz)
 *  @param q     品质因数
 *
 *  通用双二阶传递函数:
 *    H(z) = (b0 + b1*z^-1 + b2*z^-2) / (a0 + a1*z^-1 + a2*z^-2)
 *
 *  Q 值参考:
 *    0.5773  Thomson-Bessel  → 最平坦群延迟, 相位失真最小
 *    0.7071  Butterworth     → 通带最平坦, 常用默认值
 *    0.9565  Chebyshev 1dB   → 陡峭过渡带, 允许 1dB 通带波纹
 */
void biquad_init(biquad_state_t *state, biquad_type_t type, int fs, float fc, float q)
{
    float w0     = 2.0f * PI * fc / fs;
    float sin_w0 = sinf(w0);
    float cos_w0 = cosf(w0);
    float alpha  = sin_w0 / (2.0f * q);

    float b0, b1, b2, a0, a1, a2;

    switch (type) {
    case BIQUAD_LOWPASS:
        b0 = (1.0f - cos_w0) / 2.0f;
        b1 = b0 * 2.0f;
        b2 = b0;
        a0 = 1.0f + alpha;
        a1 = -2.0f * cos_w0;
        a2 = 1.0f - alpha;
        break;
    case BIQUAD_HIGHPASS:
        b0 = (1.0f + cos_w0) / 2.0f;
        b1 = -b0 * 2.0f;
        b2 = b0;
        a0 = 1.0f + alpha;
        a1 = -2.0f * cos_w0;
        a2 = 1.0f - alpha;
        break;
    case BIQUAD_BANDPASS_PEAK:
        b0 = alpha;
        b1 = 0.0f;
        b2 = -alpha;
        a0 = 1.0f + alpha;
        a1 = -2.0f * cos_w0;
        a2 = 1.0f - alpha;
        break;
    case BIQUAD_BANDSTOP_NOTCH:
        b0 = 1.0f;
        b1 = -2.0f * cos_w0;
        b2 = 1.0f;
        a0 = 1.0f + alpha;
        a1 = -2.0f * cos_w0;
        a2 = 1.0f - alpha;
        break;
    }

    // 归一化系数（除以 a0）
    state->a0 = b0 / a0;
    state->a1 = b1 / a0;
    state->a2 = b2 / a0;
    state->a3 = a1 / a0;
    state->a4 = a2 / a0;
    state->x1 = 0.0f;
    state->x2 = 0.0f;
    state->y1 = 0.0f;
    state->y2 = 0.0f;
}

/** @brief 双二阶滤波器单步迭代
 *  @param state 滤波器状态
 *  @param data  当前输入
 *  @return 当前输出
 *
 *  差分方程:
 *    y[n] = a0*x[n] + a1*x[n-1] + a2*x[n-2] - a3*y[n-1] - a4*y[n-2]
 */
float biquad_apply(biquad_state_t *state, float data)
{
    float result = state->a0 * data
                 + state->a1 * state->x1
                 + state->a2 * state->x2
                 - state->a3 * state->y1
                 - state->a4 * state->y2;
    state->x2 = state->x1;
    state->x1 = data;
    state->y2 = state->y1;
    state->y1 = result;
    return result;
}

/** @brief 快速平方根倒数 (Fast Inverse Square Root)
 *  @param x 输入正数
 *  @return 1/sqrt(x) 近似值
 *
 *  使用 Quake III 的经典 magic number 0x5f3759df
 *  精度约 0.1%, 在无 FPU 的 M0+ 上远快于 1.0f/sqrtf(x)
 */
float inv_sqrt(float x)
{
    float halfx = 0.5f * x;
    float y = x;
    long i = *(long *)&y;
    i = 0x5f3759df - (i >> 1);
    y = *(float *)&i;
    y = y * (1.5f - (halfx * y * y));   // 一次 Newton 迭代
    return y;
}