/*
 * mpu6000.c — IMU Sensor Fusion for MPU6000
 *
 * Wraps the icm42688 SPI driver (which auto-detects MPU6000) and applies
 * a complementary filter to produce a stable pitch estimate.
 *
 * Hardware: MAMBA F722S, MPU6000 on SPI1, CW270 alignment
 * Config:   Gyro ±2000°/s (init_mpu6000 sets FS_SEL=3)
 *           Accel ±16g    (init_mpu6000 sets AFS_SEL=3)
 */

#include "mpu6000.h"
#include "icm42688.h"
#include "config.h"
#include "safety.h"
#include "stm32f7xx_hal.h"
#include <math.h>

/* Scale factors matching init_mpu6000() config in icm42688.c */
#define GYRO_SCALE    (1.0f / 16.384f)   /* LSB to °/s  — ±2000°/s range */
#define ACCEL_SCALE   (1.0f / 2048.0f)   /* LSB to g    — ±16g range     */

/*
 * Complementary filter coefficient.
 * 0.98 = trust gyro integration, 0.02 = accel correction for drift.
 * Tune higher (0.99) for noisier environments.
 */
#define COMP_ALPHA    0.98f

/* Filter state */
static float s_pitch = 0.0f;
static float s_roll  = 0.0f;
static float s_yaw   = 0.0f;
static uint32_t s_last_tick = 0;

/* Gyro bias offsets (raw sensor LSBs, sensor frame) */
static float s_bias_gx   = 0.0f;
static float s_bias_gy   = 0.0f;
static float s_bias_gz   = 0.0f;
static bool  s_calibrated = false;

bool mpu6000_init(void) {
    int ret = icm42688_init();
    if (ret == 0) {
        s_pitch = 0.0f;
        s_roll  = 0.0f;
        s_yaw   = 0.0f;
        s_last_tick = HAL_GetTick();
    }
    return (ret == 0);
}

void mpu6000_calibrate(void) {
    /* LED1 + LED2 solid ON (active low) — visual indicator: calibrating */
    HAL_GPIO_WritePin(LED1_PORT, LED1_PIN, GPIO_PIN_RESET);
    HAL_GPIO_WritePin(LED2_PORT, LED2_PIN, GPIO_PIN_RESET);

    /*
     * Accumulate raw gyro readings in sensor frame.
     * Use int32 to avoid float accumulation error over 1000 samples.
     * Max raw value ~32768, 1000 samples → max sum ~32.7M < INT32_MAX.
     */
    int32_t sum_gx = 0, sum_gy = 0, sum_gz = 0;
    for (int i = 0; i < GYRO_CAL_SAMPLES; i++) {
        icm42688_data_t raw = {0};  /* zero-init: guard against icm42688_read no-op on imu_type mismatch */
        icm42688_read(&raw);
        sum_gx += raw.gx;
        sum_gy += raw.gy;
        sum_gz += raw.gz;
        HAL_Delay(1);
        /* Refresh IWDG every 40ms, starting immediately (i=0) — the gap between
         * safety_refresh() at the top of the main loop and entry here can be
         * ~10ms, so we must refresh on i=0 to avoid the 50ms IWDG window. */
        if (i % 40 == 0) safety_refresh();
    }

    s_bias_gx = (float)sum_gx / GYRO_CAL_SAMPLES;
    s_bias_gy = (float)sum_gy / GYRO_CAL_SAMPLES;
    s_bias_gz = (float)sum_gz / GYRO_CAL_SAMPLES;

    /* Reset filter state so angles start from zero with clean bias */
    s_pitch = 0.0f;
    s_roll  = 0.0f;
    s_yaw   = 0.0f;
    s_last_tick = HAL_GetTick();

    s_calibrated = true;
}

bool mpu6000_is_calibrated(void) {
    return s_calibrated;
}

void mpu6000_read(IMUData *data) {
    icm42688_data_t raw;
    icm42688_read(&raw);

    /* Compute dt from wall clock — robust to loop jitter */
    uint32_t now = HAL_GetTick();
    uint32_t elapsed_ms = now - s_last_tick;
    if (elapsed_ms == 0) elapsed_ms = 1;         /* min 1ms to avoid divide-by-zero */
    if (elapsed_ms > 100) elapsed_ms = 100;       /* clamp: don't integrate stale data */
    float dt = elapsed_ms * 0.001f;
    s_last_tick = now;

    /* Convert raw to physical units */
    float ax = raw.ax * ACCEL_SCALE;   /* g  */
    float ay = raw.ay * ACCEL_SCALE;   /* g  */
    float az = raw.az * ACCEL_SCALE;   /* g  */

    /*
     * CW270 alignment transform (Betaflight convention, R_CW270):
     *
     *   R = [[0, 1, 0], [-1, 0, 0], [0, 0, 1]]
     *   board_forward = sensor_Y    → pitch accel uses ay
     *   board_right   = -sensor_X   → roll accel uses -ax
     *   board_gx (roll rate)  =  sensor_gy   (board_gx = R[0,·]·omega_sensor)
     *   board_gy (pitch rate) = -sensor_gx   (board_gy = R[1,·]·omega_sensor)
     *   board_gz (yaw rate)   =  sensor_gz   (unchanged)
     *
     * Convention: pitch+ = nose up, roll+ = right bank, yaw+ = CW from above.
     */
    /* Subtract bias (sensor-frame raw LSBs) before axis transform + scale */
    float gyro_pitch_rate = -(raw.gx - s_bias_gx) * GYRO_SCALE;  /* °/s  board_gy = -sensor_gx */
    float gyro_roll_rate  =  (raw.gy - s_bias_gy) * GYRO_SCALE;  /* °/s  board_gx =  sensor_gy */
    float gyro_yaw_rate   =  (raw.gz - s_bias_gz) * GYRO_SCALE;  /* °/s  unchanged */

    /*
     * Accel-derived angles after CW270 transform.
     * board_ax (forward) = sensor_ay → pitch = atan2(ay, az)
     * board_ay (right)   = -sensor_ax → roll  = atan2(-ax, az)
     * Valid while total accel ≈ 1g (low linear acceleration).
     */
    float accel_pitch = atan2f( ay, az) * (180.0f / 3.14159265358979f);
    float accel_roll  = atan2f(-ax, az) * (180.0f / 3.14159265358979f);

    /*
     * Complementary filter for pitch and roll:
     *   angle = α * (angle + ω*dt) + (1−α) * accel_angle
     * Gyro integration tracks fast dynamics; accel corrects drift.
     */
    s_pitch = COMP_ALPHA * (s_pitch + gyro_pitch_rate * dt)
            + (1.0f - COMP_ALPHA) * accel_pitch;
    s_roll  = COMP_ALPHA * (s_roll  + gyro_roll_rate  * dt)
            + (1.0f - COMP_ALPHA) * accel_roll;

    /* Yaw: pure gyro integration — no accel correction, drifts over time */
    s_yaw += gyro_yaw_rate * dt;

    data->pitch      = s_pitch;
    data->pitch_rate = gyro_pitch_rate;
    data->roll       = s_roll;
    data->yaw        = s_yaw;
    data->accel_x    = ax;
    data->accel_z    = az;
}