#include "mahony_filter.h"

#include <stddef.h>

#include "fixedpoint.h"

static void mahony_filter_normalize_quaternion(mahony_filter_t *filter)
{
    const uint8_t frac_bits = filter->frac_bits;
    const int32_t norm_sq =
        FIXED_MUL(filter->q0, filter->q0, frac_bits) +
        FIXED_MUL(filter->q1, filter->q1, frac_bits) +
        FIXED_MUL(filter->q2, filter->q2, frac_bits) +
        FIXED_MUL(filter->q3, filter->q3, frac_bits);
    const int32_t inv_norm = fixed_inv_sqrt(norm_sq, frac_bits);

    if (inv_norm == 0)
    {
        mahony_filter_reset(filter);
        return;
    }

    filter->q0 = FIXED_MUL(filter->q0, inv_norm, frac_bits);
    filter->q1 = FIXED_MUL(filter->q1, inv_norm, frac_bits);
    filter->q2 = FIXED_MUL(filter->q2, inv_norm, frac_bits);
    filter->q3 = FIXED_MUL(filter->q3, inv_norm, frac_bits);
}

void mahony_filter_init(mahony_filter_t *filter, uint8_t frac_bits)
{
    if (filter == NULL)
    {
        return;
    }

    filter->frac_bits = frac_bits;
    filter->kp = 0;
    filter->ki = 0;
    mahony_filter_reset(filter);
}

void mahony_filter_reset(mahony_filter_t *filter)
{
    if (filter == NULL)
    {
        return;
    }

    filter->integral_x = 0;
    filter->integral_y = 0;
    filter->integral_z = 0;
    filter->q0 = FIXED_ONE(filter->frac_bits);
    filter->q1 = 0;
    filter->q2 = 0;
    filter->q3 = 0;
}

void mahony_filter_set_gains(mahony_filter_t *filter, int32_t kp, int32_t ki)
{
    if (filter == NULL)
    {
        return;
    }

    filter->kp = kp;
    filter->ki = ki;
}

void mahony_filter_update_imu(
    mahony_filter_t *filter,
    int32_t gx,
    int32_t gy,
    int32_t gz,
    int32_t ax,
    int32_t ay,
    int32_t az,
    int32_t dt)
{
    uint8_t frac_bits;
    int32_t half;
    int32_t accel_norm_sq;
    int32_t q0;
    int32_t q1;
    int32_t q2;
    int32_t q3;

    if (filter == NULL || dt <= 0)
    {
        return;
    }

    frac_bits = filter->frac_bits;
    half = FIXED_HALF(frac_bits);
    accel_norm_sq =
        FIXED_MUL(ax, ax, frac_bits) +
        FIXED_MUL(ay, ay, frac_bits) +
        FIXED_MUL(az, az, frac_bits);

    q0 = filter->q0;
    q1 = filter->q1;
    q2 = filter->q2;
    q3 = filter->q3;

    if (accel_norm_sq > 0)
    {
        const int32_t accel_inv_norm = fixed_inv_sqrt(accel_norm_sq, frac_bits);

        if (accel_inv_norm != 0)
        {
            const int32_t ax_n = FIXED_MUL(ax, accel_inv_norm, frac_bits);
            const int32_t ay_n = FIXED_MUL(ay, accel_inv_norm, frac_bits);
            const int32_t az_n = FIXED_MUL(az, accel_inv_norm, frac_bits);
            const int32_t vx = FIXED_MUL(FIXED_FROM_INT(2, frac_bits), (FIXED_MUL(q1, q3, frac_bits) - FIXED_MUL(q0, q2, frac_bits)), frac_bits);
            const int32_t vy = FIXED_MUL(FIXED_FROM_INT(2, frac_bits), (FIXED_MUL(q0, q1, frac_bits) + FIXED_MUL(q2, q3, frac_bits)), frac_bits);
            const int32_t vz =
                FIXED_MUL(q0, q0, frac_bits) -
                FIXED_MUL(q1, q1, frac_bits) -
                FIXED_MUL(q2, q2, frac_bits) +
                FIXED_MUL(q3, q3, frac_bits);
            const int32_t ex = FIXED_MUL(ay_n, vz, frac_bits) - FIXED_MUL(az_n, vy, frac_bits);
            const int32_t ey = FIXED_MUL(az_n, vx, frac_bits) - FIXED_MUL(ax_n, vz, frac_bits);
            const int32_t ez = FIXED_MUL(ax_n, vy, frac_bits) - FIXED_MUL(ay_n, vx, frac_bits);

            if (filter->ki != 0)
            {
                filter->integral_x += FIXED_MUL(FIXED_MUL(filter->ki, ex, frac_bits), dt, frac_bits);
                filter->integral_y += FIXED_MUL(FIXED_MUL(filter->ki, ey, frac_bits), dt, frac_bits);
                filter->integral_z += FIXED_MUL(FIXED_MUL(filter->ki, ez, frac_bits), dt, frac_bits);
            }
            else
            {
                filter->integral_x = 0;
                filter->integral_y = 0;
                filter->integral_z = 0;
            }

            gx += FIXED_MUL(filter->kp, ex, frac_bits) + filter->integral_x;
            gy += FIXED_MUL(filter->kp, ey, frac_bits) + filter->integral_y;
            gz += FIXED_MUL(filter->kp, ez, frac_bits) + filter->integral_z;
        }
    }

    {
        const int32_t half_dt = FIXED_MUL(half, dt, frac_bits);
        const int32_t q_dot0 = -FIXED_MUL(q1, gx, frac_bits) - FIXED_MUL(q2, gy, frac_bits) - FIXED_MUL(q3, gz, frac_bits);
        const int32_t q_dot1 = FIXED_MUL(q0, gx, frac_bits) + FIXED_MUL(q2, gz, frac_bits) - FIXED_MUL(q3, gy, frac_bits);
        const int32_t q_dot2 = FIXED_MUL(q0, gy, frac_bits) - FIXED_MUL(q1, gz, frac_bits) + FIXED_MUL(q3, gx, frac_bits);
        const int32_t q_dot3 = FIXED_MUL(q0, gz, frac_bits) + FIXED_MUL(q1, gy, frac_bits) - FIXED_MUL(q2, gx, frac_bits);

        filter->q0 += FIXED_MUL(q_dot0, half_dt, frac_bits);
        filter->q1 += FIXED_MUL(q_dot1, half_dt, frac_bits);
        filter->q2 += FIXED_MUL(q_dot2, half_dt, frac_bits);
        filter->q3 += FIXED_MUL(q_dot3, half_dt, frac_bits);
    }

    mahony_filter_normalize_quaternion(filter);
}

void mahony_filter_get_quaternion(
    const mahony_filter_t *filter,
    int32_t *q0,
    int32_t *q1,
    int32_t *q2,
    int32_t *q3)
{
    if (filter == NULL)
    {
        return;
    }

    if (q0 != NULL)
    {
        *q0 = filter->q0;
    }
    if (q1 != NULL)
    {
        *q1 = filter->q1;
    }
    if (q2 != NULL)
    {
        *q2 = filter->q2;
    }
    if (q3 != NULL)
    {
        *q3 = filter->q3;
    }
}

void mahony_filter_get_yaw_pitch_roll(
    const mahony_filter_t *filter,
    int32_t *yaw,
    int32_t *pitch,
    int32_t *roll)
{
    uint8_t frac_bits;
    int32_t two;
    int32_t q0q0;
    int32_t q1q1;
    int32_t q2q2;
    int32_t q3q3;
    int32_t yaw_num;
    int32_t yaw_den;
    int32_t pitch_arg;
    int32_t roll_num;
    int32_t roll_den;

    if (filter == NULL)
    {
        return;
    }

    frac_bits = filter->frac_bits;
    two = FIXED_FROM_INT(2, frac_bits);
    q0q0 = FIXED_MUL(filter->q0, filter->q0, frac_bits);
    q1q1 = FIXED_MUL(filter->q1, filter->q1, frac_bits);
    q2q2 = FIXED_MUL(filter->q2, filter->q2, frac_bits);
    q3q3 = FIXED_MUL(filter->q3, filter->q3, frac_bits);
    yaw_num = FIXED_MUL(
        two,
        FIXED_MUL(filter->q0, filter->q3, frac_bits) + FIXED_MUL(filter->q1, filter->q2, frac_bits),
        frac_bits);
    yaw_den = q0q0 + q1q1 - q2q2 - q3q3;
    pitch_arg = FIXED_MUL(
        two,
        FIXED_MUL(filter->q0, filter->q2, frac_bits) - FIXED_MUL(filter->q3, filter->q1, frac_bits),
        frac_bits);
    roll_num = FIXED_MUL(
        two,
        FIXED_MUL(filter->q0, filter->q1, frac_bits) + FIXED_MUL(filter->q2, filter->q3, frac_bits),
        frac_bits);
    roll_den = q0q0 - q1q1 - q2q2 + q3q3;

    if (yaw != NULL)
    {
        *yaw = fixed_atan2(yaw_num, yaw_den, frac_bits);
    }
    if (pitch != NULL)
    {
        *pitch = fixed_asin(pitch_arg, frac_bits);
    }
    if (roll != NULL)
    {
        *roll = fixed_atan2(roll_num, roll_den, frac_bits);
    }
}