/* This is a simple AHRS (attitude & heading reference system) based on
 * a Kalman filter.
 */

#include <ch.h>
#include <hal.h>
#include <stdio.h>
#include <stdlib.h>
#include <fix16.h>
#include <fixmatrix.h>
#include "sensor_task.h"
#include "kalman.h"
#include "kalman_model.h"
#include "ahrs_task.h"

int32_t AHRS_MAG_BIAS_X = -130;
int32_t AHRS_MAG_BIAS_Y = -60;
int32_t AHRS_MAG_BIAS_Z = -312;
int32_t AHRS_MAG_SCALE_X = 569;
int32_t AHRS_MAG_SCALE_Y = 567;
int32_t AHRS_MAG_SCALE_Z = 497;

int32_t AHRS_ACC_BIAS_X = -523;
int32_t AHRS_ACC_BIAS_Y = -335;
int32_t AHRS_ACC_BIAS_Z = -679;
int32_t AHRS_ACC_SCALE_X = 16482;
int32_t AHRS_ACC_SCALE_Y = 16496;
int32_t AHRS_ACC_SCALE_Z = 16150;

fix16_t AHRS_Q_VALUE = 2000;
fix16_t AHRS_MAG_R = 3000;
fix16_t AHRS_ACC_R = 3000;

static Thread *ahrsThread = 0;
static WORKING_AREA(ahrsThread_wa, 3500);

volatile fix16_t ahrs_north, ahrs_east, ahrs_inclination;
volatile unsigned ahrs_start_time, ahrs_updates;

static inline fix16_t deg2rad(int deg)
{
    return deg * fix16_pi / 180;
}

static inline void clamp(fix16_t *val, fix16_t min, fix16_t max)
{
    if (*val < min) *val = min;
    if (*val > max) *val = max;
}

static inline void wrap(fix16_t *val, fix16_t min, fix16_t max)
{
    while (*val > max) *val -= max - min;
    while (*val < min) *val += max - min;
}

static msg_t ahrs_task(void *arg)
{
    static EventListener el;
    
    chEvtRegister(&sensor_data_event, &el, 0);
    
    static mf16 state;
    static mf16 covariance;
    fix16_t acc_noise = 0;
    
    chRegSetThreadName("AHRS");
    
    // Initial state:
    state.rows = 5;
    state.columns = 1;
    state.errors = 0;
    state.data[0][0] = fix16_from_float(0.99);
    state.data[1][0] = state.data[2][0] = state.data[3][0] = fix16_from_float(0.1);
    state.data[4][0] = -deg2rad(70);
    
    // Note: covariance is kept in square-root form.
    covariance.rows = covariance.columns = 5;
    mf16_fill(&covariance, 0);
    covariance.data[0][0] = fix16_from_float(0.1);
    covariance.data[1][1] = fix16_from_float(0.1);
    covariance.data[2][2] = fix16_from_float(0.1);
    covariance.data[3][3] = fix16_from_float(0.1);
    covariance.data[4][4] = deg2rad(2);
    
    while (!chThdShouldTerminate())
    {
        chEvtWaitAny(ALL_EVENTS); // Wait for new sensor data
        
        // Prediction step
        mf16 Q = {5, 5, 0, {
            {AHRS_Q_VALUE, 0, 0, 0, 0},
            {0, AHRS_Q_VALUE, 0, 0, 0},
            {0, 0, AHRS_Q_VALUE, 0, 0},
            {0, 0, 0, AHRS_Q_VALUE, 0},
            {0, 0, 0, 0, 1}
        }};
        
        kalman_prediction_step(&state, &covariance, NULL, &Q);
        
        // Then perform the Kalman update step with sensor data
        {
            fix16_t w = state.data[0][0];
            fix16_t x = state.data[1][0];
            fix16_t y = state.data[2][0];
            fix16_t z = state.data[3][0];
            fix16_t sin_i = fix16_sin(state.data[4][0]);
            fix16_t cos_i = fix16_cos(state.data[4][0]);
            
            mf16 H = {1, 5};
            fix16_t Y;
            int sx, sy, sz;
            
            // 1) Copy the magnetometer data into H/Y matrices
            chSysLock();
            sx = mag_x;
            sy = mag_y;
            sz = mag_z;
            chSysUnlock();
            
            fix16_t fx = fix16_div(sx - AHRS_MAG_BIAS_X, AHRS_MAG_SCALE_X);
            fix16_t fy = fix16_div(sy - AHRS_MAG_BIAS_Y, AHRS_MAG_SCALE_Y);
            fix16_t fz = fix16_div(sz - AHRS_MAG_BIAS_Z, AHRS_MAG_SCALE_Z);
            
            // Adjust measurement noise based on how much norm differs from
            // expected magnetic field. This reduces the effect of nearby
            // magnets.
            fix16_t norm = fix16_sqrt(fix16_mul(fx, fx) + fix16_mul(fy, fy) + fix16_mul(fz, fz));
            fix16_t r_val = AHRS_MAG_R + abs(norm - fix16_one);
            
            H.data[0][0] = KALMAN_H_MAG_X_0;
            H.data[0][1] = KALMAN_H_MAG_X_1;
            H.data[0][2] = KALMAN_H_MAG_X_2;
            H.data[0][3] = KALMAN_H_MAG_X_3;
            H.data[0][4] = KALMAN_H_MAG_X_4;
            Y = fx - KALMAN_Y_MAG_X;
            kalman_potter_update(&state, &covariance, &H, r_val, Y);
            
            H.data[0][0] = KALMAN_H_MAG_Y_0;
            H.data[0][1] = KALMAN_H_MAG_Y_1;
            H.data[0][2] = KALMAN_H_MAG_Y_2;
            H.data[0][3] = KALMAN_H_MAG_Y_3;
            H.data[0][4] = KALMAN_H_MAG_Y_4;
            Y = fy - KALMAN_Y_MAG_Y;
            kalman_potter_update(&state, &covariance, &H, r_val, Y);
            
            H.data[0][0] = KALMAN_H_MAG_Z_0;
            H.data[0][1] = KALMAN_H_MAG_Z_1;
            H.data[0][2] = KALMAN_H_MAG_Z_2;
            H.data[0][3] = KALMAN_H_MAG_Z_3;
            H.data[0][4] = KALMAN_H_MAG_Z_4;
            Y = fz - KALMAN_Y_MAG_Z;
            kalman_potter_update(&state, &covariance, &H, r_val, Y);
            
            // 2) Now the accelerometer data
            chSysLock();
            sx = acc_x;
            sy = acc_y;
            sz = acc_z;
            chSysUnlock();
            
            fx = fix16_div(sx - AHRS_ACC_BIAS_X, AHRS_ACC_SCALE_X);
            fy = fix16_div(sy - AHRS_ACC_BIAS_Y, AHRS_ACC_SCALE_Y);
            fz = fix16_div(sz - AHRS_ACC_BIAS_Z, AHRS_ACC_SCALE_Z);
            
            // Adjust the measurement noise based on how much the norm of
            // measurements differs from the expected gravity. This filters
            // out any real accelerations, which could disturb the tilt
            // measurement.
            
            // The acc_noise is stored between iterations to implement an
            // IIR low-pass filter. So after a sudden movement, the next
            // measurements are down-weighted also.
            norm = fix16_sqrt(fix16_mul(fx, fx) + fix16_mul(fy, fy) + fix16_mul(fz, fz));
            acc_noise = abs(norm - fix16_one) + fix16_mul(acc_noise, fix16_from_float(0.8));
            r_val = AHRS_ACC_R + acc_noise;
            
            H.data[0][0] = KALMAN_H_ACC_X_0;
            H.data[0][1] = KALMAN_H_ACC_X_1;
            H.data[0][2] = KALMAN_H_ACC_X_2;
            H.data[0][3] = KALMAN_H_ACC_X_3;
            H.data[0][4] = KALMAN_H_ACC_X_4;
            Y = fx - KALMAN_Y_ACC_X;
            kalman_potter_update(&state, &covariance, &H, r_val, Y);
            
            H.data[0][0] = KALMAN_H_ACC_Y_0;
            H.data[0][1] = KALMAN_H_ACC_Y_1;
            H.data[0][2] = KALMAN_H_ACC_Y_2;
            H.data[0][3] = KALMAN_H_ACC_Y_3;
            H.data[0][4] = KALMAN_H_ACC_Y_4;
            Y = fy - KALMAN_Y_ACC_Y;
            kalman_potter_update(&state, &covariance, &H, r_val, Y);
            
            H.data[0][0] = KALMAN_H_ACC_Z_0;
            H.data[0][1] = KALMAN_H_ACC_Z_1;
            H.data[0][2] = KALMAN_H_ACC_Z_2;
            H.data[0][3] = KALMAN_H_ACC_Z_3;
            H.data[0][4] = KALMAN_H_ACC_Z_4;
            Y = fz - KALMAN_Y_ACC_Z;
            kalman_potter_update(&state, &covariance, &H, r_val, Y);
        }
        
        // Verify that inclination stays within +-80 degrees
        clamp(&state.data[4][0], deg2rad(-80), deg2rad(80));
        
        // And normalize the quaternion
        {
            fix16_t norm = fix16_sqrt(
                fix16_mul(state.data[0][0], state.data[0][0]) +
                fix16_mul(state.data[1][0], state.data[1][0]) +
                fix16_mul(state.data[2][0], state.data[2][0]) +
                fix16_mul(state.data[3][0], state.data[3][0])
            );
            for (int i = 0; i < 4; i++)
                state.data[i][0] = fix16_div(state.data[i][0], norm);
        }
        
        // Then rotate a Y-axis vector by the quartenion to get the
        // horizontal components of the result.
        fix16_t north, east;
        {
            fix16_t w = state.data[0][0];
            fix16_t x = state.data[1][0];
            fix16_t y = state.data[2][0];
            fix16_t z = state.data[3][0];
            east = -2 * (fix16_mul(x, y) - fix16_mul(w, z));
            north = fix16_mul(z,z) - fix16_mul(y,y) + fix16_mul(x,x) - fix16_mul(w,w);
        }
        
        fix16_t inclination = fix16_div(state.data[4][0] * 180, fix16_pi);
        chSysLock();
        ahrs_north = north;
        ahrs_east = east;
        ahrs_inclination = inclination;
        ahrs_updates++;
        chSysUnlock();
        
        if (state.errors)
        {
            printf("AHRS state errors: %x\n", state.errors);
        }
    }
    
    return 0;
}

void ahrs_start()
{
    ahrs_updates = 0;
    ahrs_start_time = chTimeNow();
    ahrsThread = chThdCreateStatic(ahrsThread_wa, sizeof(ahrsThread_wa), NORMALPRIO - 10, ahrs_task, NULL);
}

void ahrs_stop()
{
    chThdTerminate(ahrsThread);
}