#include "mpu_handler.h"
#include <stdio.h>

static float offset_ax = 0, offset_ay = 0, offset_az = 0;

static void mpu_write(uint8_t addr, uint8_t reg, uint8_t data) {
    uint8_t buf[2] = {reg, data};
    i2c_write_blocking(I2C_PORT, addr, buf, 2, false);
}

static void mpu_read(uint8_t addr, uint8_t reg, uint8_t *buf, uint8_t len) {
    i2c_write_blocking(I2C_PORT, addr, &reg, 1, true);
    i2c_read_timeout_us(I2C_PORT, addr, buf, len, false, 2500);
}

void mpu_init(void) {
    i2c_init(I2C_PORT, 400 * 1000);
    gpio_set_function(SDA_PIN, GPIO_FUNC_I2C);
    gpio_set_function(SCL_PIN, GPIO_FUNC_I2C);
    gpio_pull_up(SDA_PIN);
    gpio_pull_up(SCL_PIN);

    // Wake MPU-92.65
    mpu_write(MPU_ADDR, 0x6B, 0x00);
    sleep_ms(100);

    // Enable Bypass Mode to see the Magnetometer (AK8963)
    mpu_write(MPU_ADDR, 0x37, 0x02);
    sleep_ms(10);

    // Initialize Magnetometer: Power down then Continuous measurement mode 2 (100Hz)
    mpu_write(MAG_ADDR, 0x0A, 0x00);
    sleep_ms(10);
    mpu_write(MAG_ADDR, 0x0A, 0x16);
    sleep_ms(10);
}

void mpu_calibrate(void) {
    int num_samples = 20;
    int32_t sum_ax = 0, sum_ay = 0, sum_az = 0;
    uint8_t raw[6];

    for (int i = 0; i < num_samples; i++) {
        mpu_read(MPU_ADDR, 0x3B, raw, 6);
        sum_ax += (int16_t)((raw[0] << 8) | raw[1]);
        sum_ay += (int16_t)((raw[2] << 8) | raw[3]);
        sum_az += (int16_t)((raw[4] << 8) | raw[5]);
        sleep_ms(10);
    }

    offset_ax = (float)sum_ax / num_samples;
    offset_ay = (float)sum_ay / num_samples;
    offset_az = ((float)sum_az / num_samples) - 16384.0f;
}

void mpu_read_data(MPU_Data_t *data) {
    uint8_t raw_accel_gyro[14];
    uint8_t raw_mag[7]; // 6 data bytes + ST2 status byte

    // Read Accel (0x3B-0x40) and Gyro (0x43-0x48)
    mpu_read(MPU_ADDR, 0x3B, raw_accel_gyro, 14);

    // Process Accelerometer
    data->ax = ((float)((int16_t)((raw_accel_gyro[0] << 8) | raw_accel_gyro[1])) - offset_ax) / 16384.0f;
    data->ay = ((float)((int16_t)((raw_accel_gyro[2] << 8) | raw_accel_gyro[3])) - offset_ay) / 16384.0f;
    data->az = ((float)((int16_t)((raw_accel_gyro[4] << 8) | raw_accel_gyro[5])) - offset_az) / 16384.0f;

    // Process Gyroscope (Scale 131.0 for +/- 250 deg/s)
    data->gx = (float)((int16_t)((raw_accel_gyro[8] << 8) | raw_accel_gyro[9])) / 131.0f;
    data->gy = (float)((int16_t)((raw_accel_gyro[10] << 8) | raw_accel_gyro[11])) / 131.0f;
    data->gz = (float)((int16_t)((raw_accel_gyro[12] << 8) | raw_accel_gyro[13])) / 131.0f;

    // Read Magnetometer
    mpu_read(MAG_ADDR, 0x03, raw_mag, 7);
    // AK8963 is little-endian
    data->mx = (float)((int16_t)((raw_mag[1] << 8) | raw_mag[0]));
    data->my = (float)((int16_t)((raw_mag[3] << 8) | raw_mag[2]));
    data->mz = (float)((int16_t)((raw_mag[5] << 8) | raw_mag[4]));
}