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

// Initialize the filter
void MadgwickFilter_Init(MadgwickFilter* filter, float sampleFreq, float beta) {
    filter->q[0] = 1.0f;  // qw (scalar component)
    filter->q[1] = 0.0f;  // qx
    filter->q[2] = 0.0f;  // qy
    filter->q[3] = 0.0f;  // qz
    filter->beta = beta;  // Filter gain (e.g., 0.1)
    filter->sampleFreq = sampleFreq; // Sampling frequency in Hz
}

// Update the quaternion using accelerometer and gyroscope data
void MadgwickFilter_Update(MadgwickFilter* filter, float gx, float gy, float gz,
                          float ax, float ay, float az) {
    float q[4] = {filter->q[0], filter->q[1], filter->q[2], filter->q[3]};
    float recipNorm;
    float s[4];
    float qDot[4];
    float _2q0, _2q1, _2q2, _2q3, _4q0, _4q1, _4q2, _8q1, _8q2, q0q0, q1q1, q2q2, q3q3;

    // Convert gyroscope data from degrees/s to radians/s
    gx *= 0.0174533f; // deg/s to rad/s
    gy *= 0.0174533f;
    gz *= 0.0174533f;

    // Rate of change of quaternion from gyroscope
    qDot[0] = 0.5f * (-q[1] * gx - q[2] * gy - q[3] * gz);
    qDot[1] = 0.5f * (q[0] * gx + q[2] * gz - q[3] * gy);
    qDot[2] = 0.5f * (q[0] * gy - q[1] * gz + q[3] * gx);
    qDot[3] = 0.5f * (q[0] * gz + q[1] * gy - q[2] * gx);

    // Compute accelerometer objective function and Jacobian
    if ((ax != 0.0f) || (ay != 0.0f) || (az != 0.0f)) {
        // Normalize accelerometer measurement
        recipNorm = 1.0f / sqrtf(ax * ax + ay * ay + az * az);
        ax *= recipNorm;
        ay *= recipNorm;
        az *= recipNorm;

        // Auxiliary variables to avoid repeated calculations
        _2q0 = 2.0f * q[0];
        _2q1 = 2.0f * q[1];
        _2q2 = 2.0f * q[2];
        _2q3 = 2.0f * q[3];
        _4q0 = 4.0f * q[0];
        _4q1 = 4.0f * q[1];
        _4q2 = 4.0f * q[2];
        _8q1 = 8.0f * q[1];
        _8q2 = 8.0f * q[2];
        q0q0 = q[0] * q[0];
        q1q1 = q[1] * q[1];
        q2q2 = q[2] * q[2];
        q3q3 = q[3] * q[3];

        // Gradient descent algorithm corrective step
        s[0] = _4q0 * q2q2 + _2q2 * ax + _4q0 * q1q1 - _2q1 * ay;
        s[1] = _4q1 * q3q3 - _2q3 * ax + 4.0f * q0q0 * q[1] - _2q0 * ay - _4q1 + _8q1 * q1q1 + _8q1 * q2q2 + _4q1 * az;
        s[2] = 4.0f * q0q0 * q[2] + _2q0 * ax + _4q2 * q3q3 - _2q3 * ay - _4q2 + _8q2 * q1q1 + _8q2 * q2q2 + _4q2 * az;
        s[3] = 4.0f * q1q1 * q[3] - _2q1 * ax + 4.0f * q2q2 * q[3] - _2q2 * ay;

        // Normalize step magnitude
        recipNorm = 1.0f / sqrtf(s[0] * s[0] + s[1] * s[1] + s[2] * s[2] + s[3] * s[3]);
        s[0] *= recipNorm;
        s[1] *= recipNorm;
        s[2] *= recipNorm;
        s[3] *= recipNorm;

        // Apply feedback step
        qDot[0] -= filter->beta * s[0];
        qDot[1] -= filter->beta * s[1];
        qDot[2] -= filter->beta * s[2];
        qDot[3] -= filter->beta * s[3];
    }

    // Integrate rate of change of quaternion
    float dt = 1.0f / filter->sampleFreq;
    q[0] += qDot[0] * dt;
    q[1] += qDot[1] * dt;
    q[2] += qDot[2] * dt;
    q[3] += qDot[3] * dt;

    // Normalize quaternion
    recipNorm = 1.0f / sqrtf(q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]);
    q[0] *= recipNorm;
    q[1] *= recipNorm;
    q[2] *= recipNorm;
    q[3] *= recipNorm;

    // Update filter state
    filter->q[0] = q[0];
    filter->q[1] = q[1];
    filter->q[2] = q[2];
    filter->q[3] = q[3];
}

// Convert quaternion to Euler angles (roll, pitch, yaw) in degrees
void MadgwickFilter_GetEulerAngles(MadgwickFilter* filter, float* roll, float* pitch, float* yaw) {
    float q0 = filter->q[0];
    float q1 = filter->q[1];
    float q2 = filter->q[2];
    float q3 = filter->q[3];

    // Roll (x-axis rotation)
    *roll = atan2f(2.0f * (q0 * q1 + q2 * q3), 1.0f - 2.0f * (q1 * q1 + q2 * q2)) * 57.2958f;

    // Pitch (y-axis rotation)
    float sinp = 2.0f * (q0 * q2 - q3 * q1);
    if (fabsf(sinp) >= 1)
        *pitch = copysignf(M_PI / 2, sinp) * 57.2958f; // Use 90 degrees if out of range
    else
        *pitch = asinf(sinp) * 57.2958f;

    // Yaw (z-axis rotation) - Note: Without magnetometer, yaw will drift
    *yaw = atan2f(2.0f * (q0 * q3 + q1 * q2), 1.0f - 2.0f * (q2 * q2 + q3 * q3)) * 57.2958f;
}