// SkyLogic AeroAlign - IMU Driver Implementation
//
// MPU6050 6-axis IMU driver with complementary filter for stable angle measurement.
// Designed for static measurement (RC model setup on bench), not high-speed motion tracking.

#include "imu_driver.h"
#include "config.h"
#include <math.h>

IMU_Driver::IMU_Driver()
    : pitch_offset(0.0), roll_offset(0.0), yaw_offset(0.0),
      filtered_pitch(0.0), filtered_roll(0.0), last_update_us(0),
      alpha(COMPLEMENTARY_FILTER_ALPHA), connected(false) {
    // Initialize data structure
    memset(&data, 0, sizeof(IMU_Data));
}

bool IMU_Driver::begin(uint8_t sda_pin, uint8_t scl_pin, uint32_t i2c_freq) {
    #ifdef DEBUG_SERIAL_ENABLED
    Serial.println("[IMU] Initializing MPU6050...");
    #endif

    // Initialize I2C
    Wire.begin(sda_pin, scl_pin, i2c_freq);

    // Try to initialize MPU6050
    if (!mpu.begin(IMU_I2C_ADDRESS, &Wire)) {
        last_error = "MPU6050 not found at 0x68. Check wiring!";
        #ifdef DEBUG_SERIAL_ENABLED
        Serial.printf("[IMU] ERROR: %s\n", last_error.c_str());
        #endif
        connected = false;
        return false;
    }

    #ifdef DEBUG_SERIAL_ENABLED
    Serial.printf("[IMU] MPU6050 initialized at 0x%02X\n", IMU_I2C_ADDRESS);
    #endif

    // Configure MPU6050 settings
    // Accelerometer range: ±2g (sufficient for static measurement)
    mpu.setAccelerometerRange(MPU6050_RANGE_2_G);

    // Gyroscope range: ±250 deg/s (low range for better resolution)
    mpu.setGyroRange(MPU6050_RANGE_250_DEG);

    // Filter bandwidth: 21Hz (balance noise reduction and responsiveness)
    mpu.setFilterBandwidth(MPU6050_BAND_21_HZ);

    // Wait for IMU to stabilize
    delay(100);

    // Perform initial calibration (average first N readings)
    #ifdef DEBUG_SERIAL_ENABLED
    Serial.println("[IMU] Calibrating... (keep sensor level)");
    #endif

    float pitch_sum = 0.0;
    float roll_sum = 0.0;
    int valid_samples = 0;

    for (int i = 0; i < IMU_CALIBRATION_SAMPLES; i++) {
        sensors_event_t accel, gyro, temp;
        if (mpu.getEvent(&accel, &gyro, &temp)) {
            float pitch_raw, roll_raw;
            calculateAccelAngles(accel.acceleration.x, accel.acceleration.y, accel.acceleration.z,
                               pitch_raw, roll_raw);
            pitch_sum += pitch_raw;
            roll_sum += roll_raw;
            valid_samples++;
        }
        delay(10);  // 100Hz sampling
    }

    if (valid_samples > 0) {
        pitch_offset = pitch_sum / valid_samples;
        roll_offset = roll_sum / valid_samples;
        #ifdef DEBUG_SERIAL_ENABLED
        Serial.printf("[IMU] Calibration complete. Offsets: pitch=%.2f°, roll=%.2f°\n",
                     pitch_offset, roll_offset);
        #endif
    }

    connected = true;
    last_update_us = micros();
    return true;
}

bool IMU_Driver::update() {
    if (!connected) {
        return false;
    }

    // Get sensor events
    sensors_event_t accel, gyro, temp;
    if (!mpu.getEvent(&accel, &gyro, &temp)) {
        #ifdef DEBUG_SERIAL_ENABLED
        Serial.println("[IMU] ERROR: Failed to read sensor data");
        #endif
        return false;
    }

    // Calculate time delta (dt) in seconds
    uint32_t now_us = micros();
    float dt = (now_us - last_update_us) / 1000000.0;  // Convert to seconds
    last_update_us = now_us;

    // Prevent large dt on first update
    if (dt > 1.0 || dt <= 0.0) {
        dt = 0.01;  // Default to 10ms
    }

    // Store raw sensor data
    data.accel_x = accel.acceleration.x;
    data.accel_y = accel.acceleration.y;
    data.accel_z = accel.acceleration.z;
    data.gyro_x = gyro.gyro.x;
    data.gyro_y = gyro.gyro.y;
    data.gyro_z = gyro.gyro.z;
    data.temperature = temp.temperature;
    data.timestamp = millis();

    // Calculate pitch and roll from accelerometer (gravity vector)
    float accel_pitch, accel_roll;
    calculateAccelAngles(accel.acceleration.x, accel.acceleration.y, accel.acceleration.z,
                        accel_pitch, accel_roll);

    // Apply complementary filter (fuse gyro + accel)
    applyComplementaryFilter(accel_pitch, accel_roll, gyro.gyro.x, gyro.gyro.y, dt);

    // Apply calibration offsets
    data.pitch = constrainAngle(filtered_pitch - pitch_offset);
    data.roll = constrainAngle(filtered_roll - roll_offset);
    data.yaw = 0.0;  // Yaw not supported (requires magnetometer)

    #ifdef DEBUG_IMU_READINGS
    Serial.printf("[IMU] Pitch: %.2f°, Roll: %.2f°, Temp: %.1f°C\n",
                 data.pitch, data.roll, data.temperature);
    #endif

    return true;
}

IMU_Data IMU_Driver::getData() const {
    return data;
}

void IMU_Driver::getAngles(float &pitch, float &roll, float &yaw) const {
    pitch = data.pitch;
    roll = data.roll;
    yaw = data.yaw;
}

void IMU_Driver::calibrate() {
    #ifdef DEBUG_SERIAL_ENABLED
    Serial.println("[IMU] Calibrating offsets...");
    #endif

    // Set current angles as zero reference
    pitch_offset = filtered_pitch;
    roll_offset = filtered_roll;
    yaw_offset = 0.0;

    #ifdef DEBUG_SERIAL_ENABLED
    Serial.printf("[IMU] New offsets: pitch=%.2f°, roll=%.2f°\n",
                 pitch_offset, roll_offset);
    #endif
}

void IMU_Driver::setOffsets(float pitch_off, float roll_off, float yaw_off) {
    pitch_offset = pitch_off;
    roll_offset = roll_off;
    yaw_offset = yaw_off;

    #ifdef DEBUG_SERIAL_ENABLED
    Serial.printf("[IMU] Loaded offsets: pitch=%.2f°, roll=%.2f°, yaw=%.2f°\n",
                 pitch_offset, roll_offset, yaw_offset);
    #endif
}

void IMU_Driver::getOffsets(float &pitch_off, float &roll_off, float &yaw_off) const {
    pitch_off = pitch_offset;
    roll_off = roll_offset;
    yaw_off = yaw_offset;
}

bool IMU_Driver::isConnected() const {
    return connected;
}

String IMU_Driver::getLastError() const {
    return last_error;
}

// ========================================
// Private Methods
// ========================================

void IMU_Driver::calculateAccelAngles(float ax, float ay, float az, float &pitch, float &roll) {
    // Calculate pitch and roll from accelerometer (tilt angles)
    // Assumes sensor is stationary (accelerometer measures gravity vector)
    //
    // Pitch: Rotation around Y-axis (nose up/down)
    // Roll: Rotation around X-axis (wing tilt)
    //
    // Reference frame:
    //   X: Forward (nose direction)
    //   Y: Right wing
    //   Z: Down

    // Pitch angle (degrees)
    // atan2(ax, sqrt(ay^2 + az^2))
    pitch = atan2(ax, sqrt(ay * ay + az * az)) * 180.0 / M_PI;

    // Roll angle (degrees)
    // atan2(ay, az)
    roll = atan2(ay, az) * 180.0 / M_PI;
}

void IMU_Driver::applyComplementaryFilter(float accel_pitch, float accel_roll,
                                          float gyro_x, float gyro_y, float dt) {
    // Complementary filter: Fuse gyro (responsive) + accel (stable)
    //
    // Formula:
    //   angle = alpha * (angle + gyro * dt) + (1 - alpha) * accel_angle
    //
    // Alpha = 0.98 means:
    //   - Trust gyro 98% (fast response, but drifts over time)
    //   - Trust accel 2% (slow response, but drift-free)
    //
    // For static measurement (RC bench setup), accel dominates (no vibration).

    // Convert gyro from rad/s to deg/s
    float gyro_pitch_rate = gyro_x * 180.0 / M_PI;
    float gyro_roll_rate = gyro_y * 180.0 / M_PI;

    // Integrate gyro (predict angle change)
    float gyro_pitch = filtered_pitch + gyro_pitch_rate * dt;
    float gyro_roll = filtered_roll + gyro_roll_rate * dt;

    // Fuse gyro prediction + accel measurement
    filtered_pitch = alpha * gyro_pitch + (1.0 - alpha) * accel_pitch;
    filtered_roll = alpha * gyro_roll + (1.0 - alpha) * accel_roll;

    // Constrain to -180 to +180
    filtered_pitch = constrainAngle(filtered_pitch);
    filtered_roll = constrainAngle(filtered_roll);
}

float IMU_Driver::constrainAngle(float angle) {
    // Wrap angle to -180 to +180 range
    while (angle > 180.0) angle -= 360.0;
    while (angle < -180.0) angle += 360.0;
    return angle;
}