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

#define INV_SQRT_F32(x) (1.0f / sqrtf(x))

static I2C_HandleTypeDef *mpu_i2c;

/* ---- Mahony quaternion state ---- */
static float q0 = 1.0f, q1 = 0.0f, q2 = 0.0f, q3 = 0.0f;
static float integralFBx = 0.0f, integralFBy = 0.0f, integralFBz = 0.0f;

/* ========== Low-level I2C ========== */

static uint8_t ReadReg(uint8_t reg)
{
    uint8_t val;
    HAL_I2C_Mem_Read(mpu_i2c, MPU6050_ADDR_READ, reg, 1, &val, 1, 100);
    return val;
}

static void WriteReg(uint8_t reg, uint8_t val)
{
    HAL_I2C_Mem_Write(mpu_i2c, MPU6050_ADDR_WRITE, reg, 1, &val, 1, 100);
}

static void ReadBurst(uint8_t reg, uint8_t *buf, uint8_t len)
{
    HAL_I2C_Mem_Read(mpu_i2c, MPU6050_ADDR_READ, reg, 1, buf, len, 100);
}

/* ========== Init ========== */

uint8_t MPU6050_Init(I2C_HandleTypeDef *hi2c)
{
    mpu_i2c = hi2c;

    HAL_Delay(100);

    /* Wake up MPU6050 */
    WriteReg(MPU6050_REG_PWR_MGMT_1, 0x00);
    HAL_Delay(10);

    /* Check WHO_AM_I */
    if (ReadReg(MPU6050_REG_WHO_AM_I) != 0x68) {
        return 1;   /* error */
    }

    /* Sample rate divider: 1kHz (1 / (1 + 0)) */
    WriteReg(MPU6050_REG_SMPLRT_DIV, 0x00);

    /* DLPF: accel BW 44Hz, gyro BW 42Hz */
    WriteReg(MPU6050_REG_CONFIG, 0x03);

    /* Gyro: ±250°/s */
    WriteReg(MPU6050_REG_GYRO_CONFIG, 0x00);

    /* Accel: ±2g */
    WriteReg(MPU6050_REG_ACCEL_CONFIG, 0x00);

    return 0;
}

/* ========== Read raw ========== */

uint8_t MPU6050_ReadRaw(I2C_HandleTypeDef *hi2c, MPU6050_Data *data)
{
    uint8_t buf[14];

    if (ReadBurst(MPU6050_REG_ACCEL_XOUT_H, buf, 14) != HAL_OK)
        return 1;

    data->ax_raw = (int16_t)((buf[0]  << 8) | buf[1]);
    data->ay_raw = (int16_t)((buf[2]  << 8) | buf[3]);
    data->az_raw = (int16_t)((buf[4]  << 8) | buf[5]);

    data->gx_raw = (int16_t)((buf[8]  << 8) | buf[9]);
    data->gy_raw = (int16_t)((buf[10] << 8) | buf[11]);
    data->gz_raw = (int16_t)((buf[12] << 8) | buf[13]);

    return 0;
}

/* ========== Read scaled ========== */

uint8_t MPU6050_ReadScaled(I2C_HandleTypeDef *hi2c, MPU6050_Data *data)
{
    if (MPU6050_ReadRaw(hi2c, data) != 0)
        return 1;

    data->ax = (float)data->ax_raw / ACCEL_SCALE_2G;
    data->ay = (float)data->ay_raw / ACCEL_SCALE_2G;
    data->az = (float)data->az_raw / ACCEL_SCALE_2G;

    data->gx = (float)data->gx_raw / GYRO_SCALE_250DPS;
    data->gy = (float)data->gy_raw / GYRO_SCALE_250DPS;
    data->gz = (float)data->gz_raw / GYRO_SCALE_250DPS;

    return 0;
}

/* ========== Complementary filter ========== */

void Attitude_Complementary(Attitude *att, const MPU6050_Data *sensor, float dt, float alpha)
{
    float accel_pitch, accel_roll;

    /* Accel pitch (forward/backward tilt) */
    accel_pitch = atan2f(-sensor->ax, sqrtf(sensor->ay * sensor->ay +
                                             sensor->az * sensor->az))
                  * 57.29578f;

    /* Accel roll (left/right tilt) */
    accel_roll = atan2f(sensor->ay, sensor->az) * 57.29578f;

    /* Complementary fusion */
    att->pitch = alpha * (att->pitch + sensor->gx * dt) +
                 (1.0f - alpha) * accel_pitch;
    att->roll  = alpha * (att->roll  + sensor->gy * dt) +
                 (1.0f - alpha) * accel_roll;
}

/* ========== Mahony AHRS ========== */
/* Reference: http://x-io.co.uk/open-source-imu-and-ahrs-algorithms */

void Attitude_MahonyUpdate(Attitude *att, const MPU6050_Data *sensor, float dt)
{
    float recipNorm;
    float halfvx, halfvy, halfvz;
    float halfex, halfey, halfez;
    float qa, qb, qc;

    float ax = sensor->ax, ay = sensor->ay, az = sensor->az;
    float gx = sensor->gx, gy = sensor->gy, gz = sensor->gz;

    /* Normalize accel */
    recipNorm = INV_SQRT_F32(ax * ax + ay * ay + az * az);
    ax *= recipNorm; ay *= recipNorm; az *= recipNorm;

    /* Estimated gravity from quaternion */
    halfvx = q1 * q3 - q0 * q2;
    halfvy = q0 * q1 + q2 * q3;
    halfvz = q0 * q0 - 0.5f + q3 * q3;

    /* Cross product error */
    halfex = (ay * halfvz - az * halfvy);
    halfey = (az * halfvx - ax * halfvz);
    halfez = (ax * halfvy - ay * halfvx);

    /* PI controller */
    integralFBx += MAHONY_KI * halfex * dt;
    integralFBy += MAHONY_KI * halfey * dt;
    integralFBz += MAHONY_KI * halfez * dt;

    gx += integralFBx + MAHONY_KP * halfex;
    gy += integralFBy + MAHONY_KP * halfey;
    gz += integralFBz + MAHONY_KP * halfez;

    /* Integrate rate of change of quaternion (1st order RK) */
    qa = q0; qb = q1; qc = q2;
    q0 += (-qb * gx - qc * gy - q3 * gz) * 0.5f * dt;
    q1 += ( qa * gx + qc * gz - q3 * gy) * 0.5f * dt;
    q2 += ( qa * gy - qb * gz + q3 * gx) * 0.5f * dt;
    q3 += ( qa * gz + qb * gy - qc * gx) * 0.5f * dt;

    /* Normalize quaternion */
    recipNorm = INV_SQRT_F32(q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3);
    q0 *= recipNorm; q1 *= recipNorm; q2 *= recipNorm; q3 *= recipNorm;

    /* Extract pitch and roll */
    att->pitch = asinf(-2.0f * q1 * q3 + 2.0f * q0 * q2) * 57.29578f;
    att->roll  = atan2f(2.0f * q2 * q3 + 2.0f * q0 * q1,
                        -2.0f * q1 * q1 - 2.0f * q2 * q2 + 1.0f) * 57.29578f;
}

/* ========== Gyro calibration ========== */

void MPU6050_CalibrateGyro(I2C_HandleTypeDef *hi2c,
                           float *offset_x, float *offset_y, float *offset_z,
                           uint16_t samples)
{
    float sum_x = 0, sum_y = 0, sum_z = 0;
    MPU6050_Data data;

    for (uint16_t i = 0; i < samples; i++) {
        MPU6050_ReadScaled(hi2c, &data);
        sum_x += data.gx;
        sum_y += data.gy;
        sum_z += data.gz;
        HAL_Delay(1);
    }

    *offset_x = sum_x / samples;
    *offset_y = sum_y / samples;
    *offset_z = sum_z / samples;
}