/* This is a simple AHRS (attitude & heading reference system) based on
 * a Kalman filter.
 * 
 * The coordinate system is tied to the platform inside the sphere, and
 * the orientation of the motor axis. The rotations are in the following
 * order:
 * 1) Heading: rotation of motor axis around vertical axis. When axis is
 *    perpendicular to north-south direction, angle is 0.
 * 2) Roll: rotation of motor axis around horizontal axis. When axis is
 *    horizontal the angle is 0.
 * 3) Pitch: rotation around the motor axis. When counterweight is pointing
 *    down the angle is 0.
 */

#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"
#include "control.h"
#include "motor.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_MAG_SERVO_BIAS_X_A = -4206495;
fix16_t AHRS_MAG_SERVO_BIAS_X_B = 3833011;
fix16_t AHRS_MAG_SERVO_BIAS_X_C = 8948018;
fix16_t AHRS_MAG_SERVO_BIAS_Y_A = -168014;
fix16_t AHRS_MAG_SERVO_BIAS_Y_B = -3121362;
fix16_t AHRS_MAG_SERVO_BIAS_Y_C = -2578856;
fix16_t AHRS_MAG_SERVO_BIAS_Z_A = -2749007;
fix16_t AHRS_MAG_SERVO_BIAS_Z_B = 7006532;
fix16_t AHRS_MAG_SERVO_BIAS_Z_C = 10005582;

void get_calibrated_mag(fix16_t *fx, fix16_t *fy, fix16_t *fz, fix16_t *norm)
{
    chSysLock();
    int x = mag_x;
    int y = mag_y;
    int z = mag_z;
    fix16_t s = servo_position;
    chSysUnlock();
    
    fix16_t s2 = fix16_mul(s, s);
    int bx = fix16_to_int(fix16_mul(AHRS_MAG_SERVO_BIAS_X_A, s2)
                        + fix16_mul(AHRS_MAG_SERVO_BIAS_X_B, s)
                        +           AHRS_MAG_SERVO_BIAS_X_C);
    int by = fix16_to_int(fix16_mul(AHRS_MAG_SERVO_BIAS_Y_A, s2)
                        + fix16_mul(AHRS_MAG_SERVO_BIAS_Y_B, s)
                        +           AHRS_MAG_SERVO_BIAS_Y_C);
    int bz = fix16_to_int(fix16_mul(AHRS_MAG_SERVO_BIAS_Z_A, s2)
                        + fix16_mul(AHRS_MAG_SERVO_BIAS_Z_B, s)
                        +           AHRS_MAG_SERVO_BIAS_Z_C);
    
    *fx = fix16_div(x - AHRS_MAG_BIAS_X - bx, AHRS_MAG_SCALE_X);
    *fy = fix16_div(y - AHRS_MAG_BIAS_Y - by, AHRS_MAG_SCALE_Y);
    *fz = fix16_div(z - AHRS_MAG_BIAS_Z - bz, AHRS_MAG_SCALE_Z);
    
    *norm = fix16_sqrt(fix16_mul(*fx, *fx) + fix16_mul(*fy, *fy) + fix16_mul(*fz, *fz));
}

void get_calibrated_acc(fix16_t *fx, fix16_t *fy, fix16_t *fz, fix16_t *norm)
{
    chSysLock();
    int x = acc_x;
    int y = acc_y;
    int z = acc_z;
    chSysUnlock();
    
    *fx = fix16_div(x - AHRS_ACC_BIAS_X, AHRS_ACC_SCALE_X);
    *fy = fix16_div(y - AHRS_ACC_BIAS_Y, AHRS_ACC_SCALE_Y);
    *fz = fix16_div(z - AHRS_ACC_BIAS_Z, AHRS_ACC_SCALE_Z);
    *norm = fix16_sqrt(fix16_mul(*fx, *fx) + fix16_mul(*fy, *fy) + fix16_mul(*fz, *fz));
}

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

volatile fix16_t ahrs_heading, ahrs_roll, ahrs_pitch, 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:
    // heading = 0° +- 180°
    // roll = 0° +- 90°
    // pitch = 0° +- 90°
    // magnetic inclination = 70° +- 10°
    state.rows = 4;
    state.columns = 1;
    state.errors = 0;
    state.data[0][0] = state.data[1][0] = state.data[2][0] = 0;
    state.data[3][0] = -deg2rad(70);
    
    // Note: covariance is kept in square-root form.
    covariance.rows = covariance.columns = 4;
    mf16_fill(&covariance, 0);
    covariance.data[0][0] = deg2rad(180);
    covariance.data[1][1] = deg2rad(90);
    covariance.data[2][2] = deg2rad(90);
    covariance.data[3][3] = deg2rad(10);
    
    while (!chThdShouldTerminate())
    {
        chEvtWaitAny(ALL_EVENTS); // Wait for new sensor data
        
        // Perform kalman prediction step (assumed 200Hz freq)
        mf16 Q = {4, 4, 0, {
            {500, 0, 0, 0}, // 90 deg/s
            {0, 2000, 0, 0}, // 360 deg/s
            {0, 0, 2000, 0},
            {0, 0, 0, 1} // 0.2 deg/s
        }};
        
        kalman_prediction_step(&state, &covariance, NULL, &Q);
        
        // Then perform the Kalman update step with sensor data
        {
            fix16_t sin_h = fix16_sin(state.data[0][0]);
            fix16_t cos_h = fix16_cos(state.data[0][0]);
            fix16_t sin_r = fix16_sin(state.data[1][0]);
            fix16_t cos_r = fix16_cos(state.data[1][0]);
            fix16_t sin_p = fix16_sin(state.data[2][0]);
            fix16_t cos_p = fix16_cos(state.data[2][0]);
            fix16_t sin_i = fix16_sin(state.data[3][0]);
            fix16_t cos_i = fix16_cos(state.data[3][0]);
            
            mf16 H = {6, 4};
            mf16 Y = {6, 1};
            mf16 R = {6, 6};
            
            // 1) Copy the magnetometer data into H/Y matrices
            fix16_t fx, fy, fz, norm;
            get_calibrated_mag(&fx, &fy, &fz, &norm);
            
            // Adjust measurement noise based on how much norm differs from
            // expected magnetic field. This reduces the effect of nearby
            // magnets.
            fix16_t r_val = fix16_from_float(0.05) + abs(norm - fix16_one);
            R.data[0][0] = R.data[1][1] = R.data[2][2] = r_val;
            
            Y.data[0][0] = fx - KALMAN_Y_MAG_X;
            Y.data[1][0] = fy - KALMAN_Y_MAG_Y;
            Y.data[2][0] = fz - KALMAN_Y_MAG_Z;
            
            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[1][0] = KALMAN_H_MAG_Y_0;
            H.data[1][1] = KALMAN_H_MAG_Y_1;
            H.data[1][2] = KALMAN_H_MAG_Y_2;
            H.data[1][3] = KALMAN_H_MAG_Y_3;
            
            H.data[2][0] = KALMAN_H_MAG_Z_0;
            H.data[2][1] = KALMAN_H_MAG_Z_1;
            H.data[2][2] = KALMAN_H_MAG_Z_2;
            H.data[2][3] = KALMAN_H_MAG_Z_3;
         
            // 2) Now the accelerometer data
            get_calibrated_acc(&fx, &fy, &fz, &norm);
            
            // 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.
            acc_noise = abs(norm - fix16_one) + fix16_mul(acc_noise, fix16_from_float(0.8));
            r_val = fix16_from_float(0.05) + acc_noise;
            R.data[3][3] = R.data[4][4] = R.data[5][5] = r_val;
            
            Y.data[3][0] = fx - KALMAN_Y_ACC_X;
            Y.data[4][0] = fy - KALMAN_Y_ACC_Y;
            Y.data[5][0] = fz - KALMAN_Y_ACC_Z;
            
            H.data[3][0] = KALMAN_H_ACC_X_0;
            H.data[3][1] = KALMAN_H_ACC_X_1;
            H.data[3][2] = KALMAN_H_ACC_X_2;
            H.data[3][3] = KALMAN_H_ACC_X_3;
            
            H.data[4][0] = KALMAN_H_ACC_Y_0;
            H.data[4][1] = KALMAN_H_ACC_Y_1;
            H.data[4][2] = KALMAN_H_ACC_Y_2;
            H.data[4][3] = KALMAN_H_ACC_Y_3;
            
            H.data[5][0] = KALMAN_H_ACC_Z_0;
            H.data[5][1] = KALMAN_H_ACC_Z_1;
            H.data[5][2] = KALMAN_H_ACC_Z_2;
            H.data[5][3] = KALMAN_H_ACC_Z_3;
            
            // Then perform the actual computations
            kalman_update(&state, &covariance, &H, &R, &Y);
        }
        
        // Verify that inclination stays within +-80 degrees
        clamp(&state.data[3][0], deg2rad(-80), deg2rad(80));
        
        // And roll within +- 85 degrees
        clamp(&state.data[1][0], deg2rad(-85), deg2rad(85));
        
        // And wrap heading and pitch if they make a full circle
        wrap(&state.data[0][0], 0, 2 * fix16_pi);
        wrap(&state.data[2][0], -fix16_pi, fix16_pi);
        
        fix16_t scale = fix16_div(fix16_from_int(180), fix16_pi);
        chSysLock();
        ahrs_heading = fix16_mul(state.data[0][0], scale);
        ahrs_roll = fix16_mul(state.data[1][0], scale);
        ahrs_pitch = fix16_mul(state.data[2][0], scale) - fix16_from_int(10);
        ahrs_inclination = fix16_mul(state.data[3][0], scale);
        ahrs_updates++;
        chSysUnlock();
        
        run_control();
    }
    
    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);
}