//#DOC Main, program entry and main loop
#include "main.h"
#include <otSerial.h>
#include <otPrintf.h>
#include <otTimer.h>
#include <otDisplay.h>
#include <otILI9341.h>
#include <otSpi.h>
#include <otI2cBus.h>
#include <otFPGA.h>
#include <otMatrix.h>
#include <otEKF.h>
#include <otMPU9250.h>
#include <otTTY.h>
#include <otMath.h>
#include <otMalloc.h>


#define DTYPE otDisplay::D_7SEGMENT
//#define DTYPE otDisplay::D_NIXIE

// State Space dimension
#define SS_X_LEN    (4)
#define SS_Z_LEN    (6)
#define SS_U_LEN    (3)
#define SS_DT_MILIS (20)                        // 20 ms
#define SS_DT       F32(SS_DT_MILIS / 1000.)    // Sampling time

#define M_USAGE printf("%4d: Stack = %5d, Mem. = %6d\n", __LINE__, _stackUsage(), _memUsage - _memFree());

// Nonlinear & linearization function --------------------------------------------------------------------------------------
bool Main_bUpdateNonlinearX(otMatrix & X_Next, otMatrix & X, otMatrix & U);
bool Main_bUpdateNonlinearY(otMatrix & Y, otMatrix & X, otMatrix & U);
bool Main_bCalcJacobianF(otMatrix & F, otMatrix & X, otMatrix & U);
bool Main_bCalcJacobianH(otMatrix & H, otMatrix & X, otMatrix & U);

U32     _memUsage;

otDisplay       dsp_y;
otDisplay       dsp_p;
otDisplay       dsp_r;
otILI9341       tft;
otTTY           tty;
otMPU9250       imu;

// ================================================= EKF Variables/function declaration =================================================
// EKF initialization constant
#define P_INIT      (10.)
#define Q_INIT      (1e-6)
#define R_INIT_ACC  (0.0015 / 10.)
#define R_INIT_MAG  (0.0015 / 10.)
// P(k=0) variable ---------------------------------------------------------------------------------------------------------
F32 EKF_PINIT_data[SS_X_LEN * SS_X_LEN] = {     P_INIT, 0,      0,      0,
                                                0,      P_INIT, 0,      0,
                                                0,      0,      P_INIT, 0,
                                                0,      0,      0,      P_INIT};
otMatrix EKF_PINIT;

// Q constant --------------------------------------------------------------------------------------------------------------
F32 EKF_QINIT_data[SS_X_LEN * SS_X_LEN] = {     Q_INIT, 0,      0,      0,
                                                0,      Q_INIT, 0,      0,
                                                0,      0,      Q_INIT, 0,
                                                0,      0,      0,      Q_INIT};
                                                
otMatrix EKF_QINIT;

// R constant --------------------------------------------------------------------------------------------------------------
F32 EKF_RINIT_data[SS_Z_LEN * SS_Z_LEN] = {     R_INIT_ACC, 0,          0,          0,          0,          0,
                                                0,          R_INIT_ACC, 0,          0,          0,          0,
                                                0,          0,          R_INIT_ACC, 0,          0,          0,
                                                0,          0,          0,          R_INIT_MAG, 0,          0,
                                                0,          0,          0,          0,          R_INIT_MAG, 0,
                                                0,          0,          0,          0,          0,          R_INIT_MAG};
otMatrix EKF_RINIT;

// EKF variables -----------------------------------------------------------------------------------------------------------
otMatrix    quaternionData;
otMatrix    Y;
otMatrix    U;
otEKF       EKF_IMU;

// =============================================== Sharing Variables/function declaration ===============================================
// Gravity vector constant (align with global Z-axis)
#define IMU_ACC_Z0          (1)

// Magnetic vector constant (align with local magnetic vector)
F32 IMU_MAG_B0_data[3] = { (F32) cos(0), (F32) sin(0), 0.000000 };
otMatrix IMU_MAG_B0;

// The hard-magnet bias
F32 HARD_IRON_BIAS_data[3] = { 8.832973, 7.243323, 23.95714 };
otMatrix HARD_IRON_BIAS;

F32 yaw, pitch, roll;

// IMU calibrarion and selftest
void imuCalibrate()
{
    F32   SelfTest[6];  // Holds results of gyro and accelerometer self test
    imu.reset();
    tty.print("IMU Self test started ...\n");
    imu.selfTest(SelfTest);
    tty.print("Acc. = X:%6.1f, Y:%6.1f, Z:%6.1f %% of FV\n", SelfTest[0], SelfTest[1], SelfTest[2]);
    tty.print("Gyr. = X:%6.1f, Y:%6.1f, Z:%6.1f %% of FV\n", SelfTest[3], SelfTest[4], SelfTest[5]);
    tty.print("Calibrating gyroscope and accelerometer ...\n");
    imu.accelGyroCalibration(); // Calibrate gyro and accelerometers, load biases in bias registers
    tty.print("Acc.(mg) = X:%6.1f, Y:%6.1f, Z:%6.1f\n", 1000. * imu.d->accelBias[0], 1000. * imu.d->accelBias[1], 1000. * imu.d->accelBias[2]);
    tty.print("Gyr.(dps)= X:%6.1f, Y:%6.1f, Z:%6.1f\n", imu.d->gyroBias[0], imu.d->gyroBias[1], imu.d->gyroBias[2]);
    imu.setRange(otMPU9250::AFS_2G, otMPU9250::GFS_250DPS, otMPU9250::MFS_16BITS);
    imu.imuInit();
    if(imu.magnetometerCheck())
    {
        // Get magnetometer calibration from AK8963 ROM
        imu.magnetometerInit();
        tty.print("Magnetometer initialized for active data mode ...\n"); // Initialize device for active mode read of magnetometer
        imu.magnetometerCalibration();
        tty.print("Mag. (mG)  = X:%6.1f, Y:%6.1f, Z:%6.1f\n", imu.d->magBias[0], imu.d->magBias[1], imu.d->magBias[2]);
        tty.print("Mag. (mG)  = X:%6.1f, Y:%6.1f, Z:%6.1f\n", imu.d->magScale[0], imu.d->magScale[1], imu.d->magScale[2]);
    }
}

void taitBryanAngles(F32 & yaw, F32 & pitch, F32 & roll)
{
    yaw   = atan2(2.0f * (quaternionData[1][0] * quaternionData[2][0] + quaternionData[0][0] * quaternionData[3][0]), quaternionData[0][0] * quaternionData[0][0] + quaternionData[1][0] * quaternionData[1][0] - quaternionData[2][0] * quaternionData[2][0] - quaternionData[3][0] * quaternionData[3][0]);   
    pitch = -asin(2.0f * (quaternionData[1][0] * quaternionData[3][0] - quaternionData[0][0] * quaternionData[2][0]));
    roll  = atan2(2.0f * (quaternionData[0][0] * quaternionData[1][0] + quaternionData[2][0] * quaternionData[3][0]), quaternionData[0][0] * quaternionData[0][0] - quaternionData[1][0] * quaternionData[1][0] - quaternionData[2][0] * quaternionData[2][0] + quaternionData[3][0] * quaternionData[3][0]);
    pitch *= 180.0f / M_PI;
    yaw   *= 180.0f / M_PI;
    roll  *= 180.0f / M_PI;
}

void setup()
{
    // Set stdout
    UartEnableSignals(UART0);
    UartSetBaud(UART0, 115200);
    setStdout(1024, UART0); // printf output is now directed on UART0
    
    // For process sync
    TimerSetAsCounter(TIMER0);
    TimerEnable(TIMER0, true);
    
    // Init low level
    FpgaInit();
    FpgaHiddenPort(HP_EXT_I2C, true);
    
    I2cBusReset();
    I2cBusFrequency(400, 30);
    
    // Debug memory leakage
    _memUsage = _memFree();
    
    // Init display
    tft.init();
    tft.setRotation(otILI9341::IR_270);
    tft.clear();
    tft.setTextColor(RGB565_YELLOW, RGB565_BLACK);
    
    tty.init(1, &tft);
    tty.print("SYSTEM INIT, PLEASE WAIT ...\n");
    
    // Init IMU
    imu.init();
    if(imu.check())
        imuCalibrate();
    else
    {
        tty.print("Can't find MPU9250\n");
        IDLE;
    }
    
    // EKF working variables
    EKF_PINIT.create(SS_X_LEN, SS_X_LEN, EKF_PINIT_data);
    EKF_QINIT.create(SS_X_LEN, SS_X_LEN, EKF_QINIT_data);
    EKF_RINIT.create(SS_Z_LEN, SS_Z_LEN, EKF_RINIT_data);
    
    quaternionData.create(SS_X_LEN, 1);
    Y.create(SS_Z_LEN, 1);
    U.create(SS_U_LEN, 1);
    
    IMU_MAG_B0.create(3, 1, IMU_MAG_B0_data);
    HARD_IRON_BIAS.create(3, 1, HARD_IRON_BIAS_data);
    
    // EKF initialization -----------------------------------------
    // x(k=0) = [1 0 0 0]'
    quaternionData.vSetToZero();
    quaternionData[0][0] = 1.0;
    EKF_IMU.create(quaternionData, EKF_PINIT, EKF_QINIT, EKF_RINIT, Main_bUpdateNonlinearX, Main_bUpdateNonlinearY, Main_bCalcJacobianF, Main_bCalcJacobianH);
 
    // Seven segment display init
    dsp_y.init(DTYPE, 3, "%6.1f", &tft);
    dsp_p.init(DTYPE, 3, "%6.1f", &tft);
    dsp_r.init(DTYPE, 3, "%6.1f", &tft);
    
    tft.clear();
    // Some text on display ....
    tft.setTextSize(3);
    tft.setCursor(5, 20);
    tft.writeString("Yaw");
    tft.setCursor(5, 95);
    tft.writeString("Pitch");
    tft.setCursor(5, 165);
    tft.writeString("Roll");
    tft.setTextSize(1);
    
    M_USAGE
}

void loop()
{
    F32     Ax, Ay, Az;
    F32     Bx, By, Bz;
    F32     p, q, r;
    S16     pos = (320 - dsp_y.horSize()) / 2;
    while(1)
    {
        if(ElapsedTimeCheck(SS_DT_MILIS, TIMER0))
        {
            ElapsedTimeStart(TIMER0);
            
            //M_USAGE

            imu.accelerometer(Ax, Ay, Az);
            imu.gyroscope(p, q, r);
            imu.magnetometer(Bx, By, Bz);

            p = DEG_TO_RAD(p);
            q = DEG_TO_RAD(q);
            r = DEG_TO_RAD(r);
            
            // ================== Read the sensor data / simulate the system here ==================
            
            // Input 1:3 = gyroscope
            U[0][0] = p;  U[1][0] = q;  U[2][0] = r;
            // Output 1:3 = accelerometer
            Y[0][0] = Ax; Y[1][0] = Ay; Y[2][0] = Az;
            // Output 4:6 = magnetometer
            Y[3][0] = Bx; Y[4][0] = By; Y[5][0] = Bz;

            // Compensating Hard-Iron Bias for magnetometer
            Y[3][0] = Y[3][0] - HARD_IRON_BIAS[0][0];
            Y[4][0] = Y[4][0] - HARD_IRON_BIAS[1][0];
            Y[5][0] = Y[5][0] - HARD_IRON_BIAS[2][0];

            // Normalizing the output vector */
            F32 _normG = sqrt(Y[0][0] * Y[0][0]) + (Y[1][0] * Y[1][0]) + (Y[2][0] * Y[2][0]);
            Y[0][0] = Y[0][0] / _normG;
            Y[1][0] = Y[1][0] / _normG;
            Y[2][0] = Y[2][0] / _normG;
            F32 _normM = sqrt(Y[3][0] * Y[3][0]) + (Y[4][0] * Y[4][0]) + (Y[5][0] * Y[5][0]);
            Y[3][0] = Y[3][0] / _normM;
            Y[4][0] = Y[4][0] / _normM;
            Y[5][0] = Y[5][0] / _normM;
            // ------------------ Read the sensor data / simulate the system here ------------------

            // ============================= Update the Kalman Filter ==============================
            if(!EKF_IMU.bUpdate(Y, U))
            {
                quaternionData.vSetToZero();
                quaternionData[0][0] = 1.0;
                EKF_IMU.vReset(quaternionData, EKF_PINIT, EKF_QINIT, EKF_RINIT);
                printf("Error!\n");
            }
            //M_USAGE
            
            quaternionData = EKF_IMU.GetX();
            taitBryanAngles(yaw, pitch, roll);
            
            dsp_y.display(yaw, pos + 40, 10);
            dsp_p.display(pitch, pos + 40, 80);
            dsp_r.display(roll, pos + 40, 150);
        }
    }
}

bool Main_bUpdateNonlinearX(otMatrix & X_Next, otMatrix & X, otMatrix & U)
{
    /* Insert the nonlinear update transformation here
     *          x(k+1) = f[x(k), u(k)]
     *
     * The quaternion update function:
     *  q0_dot = 1/2. * (  0   - p*q1 - q*q2 - r*q3)
     *  q1_dot = 1/2. * ( p*q0 +   0  + r*q2 - q*q3)
     *  q2_dot = 1/2. * ( q*q0 - r*q1 +  0   + p*q3)
     *  q3_dot = 1/2. * ( r*q0 + q*q1 - p*q2 +  0  )
     *
     * Euler method for integration:
     *  q0 = q0 + q0_dot * dT;
     *  q1 = q1 + q1_dot * dT;
     *  q2 = q2 + q2_dot * dT;
     *  q3 = q3 + q3_dot * dT;
     */
    F32     q0, q1, q2, q3;
    F32     p, q, r;
    
    q0 = X[0][0];
    q1 = X[1][0];
    q2 = X[2][0];
    q3 = X[3][0];
    
    p = U[0][0];
    q = U[1][0];
    r = U[2][0];
    
    X_Next[0][0] = (0.5 * (+0.0 - p * q1 - q * q2 - r * q3)) * SS_DT + q0;
    X_Next[1][0] = (0.5 * (+ p * q0 +0.0 + r * q2 -q * q3)) * SS_DT + q1;
    X_Next[2][0] = (0.5 * (+ q * q0 - r * q1 +0.0 + p * q3)) * SS_DT + q2;
    X_Next[3][0] = (0.5 * (+ r * q0 + q * q1 - p * q2 +0.0)) * SS_DT + q3;
    
    
    // ======= Additional ad-hoc quaternion normalization to make sure the quaternion is a unit vector (i.e. ||q|| = 1) =======
    if(!X_Next.bNormVector())
        // System error, return false so we can reset the UKF
        return false;
    
    return true;
}

bool Main_bUpdateNonlinearY(otMatrix & Y, otMatrix & X, otMatrix & U)
{
    /* Insert the nonlinear measurement transformation here
     *          y(k)   = h[x(k), u(k)]
     *
     * The measurement output is the gravitational and magnetic projection to the body
     *     DCM     = [(+(q0**2)+(q1**2)-(q2**2)-(q3**2)),                        2*(q1*q2+q0*q3),                        2*(q1*q3-q0*q2)]
     *               [                   2*(q1*q2-q0*q3),     (+(q0**2)-(q1**2)+(q2**2)-(q3**2)),                        2*(q2*q3+q0*q1)]
     *               [                   2*(q1*q3+q0*q2),                        2*(q2*q3-q0*q1),     (+(q0**2)-(q1**2)-(q2**2)+(q3**2))]
     *
     *  G_proj_sens = DCM * [0 0 1]             --> Gravitational projection to the accelerometer sensor
     *  M_proj_sens = DCM * [Mx My Mz]          --> (Earth) magnetic projection to the magnetometer sensor
     */
    F32     q0, q1, q2, q3;
    F32     q0_2, q1_2, q2_2, q3_2;

    q0 = X[0][0];
    q1 = X[1][0];
    q2 = X[2][0];
    q3 = X[3][0];

    q0_2 = q0 * q0;
    q1_2 = q1 * q1;
    q2_2 = q2 * q2;
    q3_2 = q3 * q3;
    
    Y[0][0] = (2 * q1 * q3 -2 * q0 * q2) * IMU_ACC_Z0;

    Y[1][0] = (2 * q2 * q3 +2 * q0 * q1) * IMU_ACC_Z0;

    Y[2][0] = (+(q0_2) -(q1_2) -(q2_2) +(q3_2)) * IMU_ACC_Z0;
    
    Y[3][0] = (+(q0_2) + (q1_2) - (q2_2) - (q3_2)) * IMU_MAG_B0[0][0] + (2 * (q1 * q2 + q0 * q3)) * IMU_MAG_B0[1][0] + (2 * (q1 * q3 - q0 * q2)) * IMU_MAG_B0[2][0];

    Y[4][0] = (2 * (q1 * q2 - q0 * q3)) * IMU_MAG_B0[0][0] + (+(q0_2) - (q1_2) + (q2_2) - (q3_2)) * IMU_MAG_B0[1][0] + (2 * (q2 * q3 + q0 * q1)) * IMU_MAG_B0[2][0];

    Y[5][0] = (2 * (q1 * q3 + q0 * q2)) * IMU_MAG_B0[0][0] + (2 * (q2 * q3 - q0 * q1)) * IMU_MAG_B0[1][0] + (+(q0_2) - (q1_2) - (q2_2) + (q3_2)) * IMU_MAG_B0[2][0];
    return true;
}

bool Main_bCalcJacobianF(otMatrix & F, otMatrix & X, otMatrix & U)
{
    /* In Main_bUpdateNonlinearX():
     *  q0 = q0 + q0_dot * dT;
     *  q1 = q1 + q1_dot * dT;
     *  q2 = q2 + q2_dot * dT;
     *  q3 = q3 + q3_dot * dT;
     */
    F32     p, q, r;

    p = U[0][0];
    q = U[1][0];
    r = U[2][0];

    F[0][0] =  1.000;
    F[1][0] =  0.5 * p * SS_DT;
    F[2][0] =  0.5 * q * SS_DT;
    F[3][0] =  0.5 * r * SS_DT;

    F[0][1] = -0.5 * p * SS_DT;
    F[1][1] =  1.000;
    F[2][1] = -0.5 * r * SS_DT;
    F[3][1] =  0.5 * q * SS_DT;

    F[0][2] = -0.5 * q * SS_DT;
    F[1][2] =  0.5 * r * SS_DT;
    F[2][2] =  1.000;
    F[3][2] = -0.5 * p * SS_DT;

    F[0][3] = -0.5 * r * SS_DT;
    F[1][3] = -0.5 * q * SS_DT;
    F[2][3] =  0.5 * p * SS_DT;
    F[3][3] =  1.000;
    
    return true;
}

bool Main_bCalcJacobianH(otMatrix & H, otMatrix & X, otMatrix & U)
{
    F32     q0, q1, q2, q3;

    q0 = X[0][0];
    q1 = X[1][0];
    q2 = X[2][0];
    q3 = X[3][0];
    
    H[0][0] = -2 * q2 * IMU_ACC_Z0;
    H[1][0] = +2 * q1 * IMU_ACC_Z0;
    H[2][0] = +2 * q0 * IMU_ACC_Z0;
    H[3][0] =  2 * q0 * IMU_MAG_B0[0][0] + 2 * q3 * IMU_MAG_B0[1][0] - 2 * q2 * IMU_MAG_B0[2][0];
    H[4][0] = -2 * q3 * IMU_MAG_B0[0][0] + 2 * q0 * IMU_MAG_B0[1][0] + 2 * q1 * IMU_MAG_B0[2][0];
    H[5][0] =  2 * q2 * IMU_MAG_B0[0][0] - 2 * q1 * IMU_MAG_B0[1][0] + 2 * q0 * IMU_MAG_B0[2][0];
    
    H[0][1] = +2 * q3 * IMU_ACC_Z0;
    H[1][1] = +2 * q0 * IMU_ACC_Z0;
    H[2][1] = -2 * q1 * IMU_ACC_Z0;
    H[3][1] =  2 * q1 * IMU_MAG_B0[0][0] + 2 * q2 * IMU_MAG_B0[1][0] + 2 * q3 * IMU_MAG_B0[2][0];
    H[4][1] =  2 * q2 * IMU_MAG_B0[0][0] - 2 * q1 * IMU_MAG_B0[1][0] + 2 * q0 * IMU_MAG_B0[2][0];
    H[5][1] =  2 * q3 * IMU_MAG_B0[0][0] - 2 * q0 * IMU_MAG_B0[1][0] - 2 * q1 * IMU_MAG_B0[2][0];
    
    H[0][2] = -2 * q0 * IMU_ACC_Z0;
    H[1][2] = +2 * q3 * IMU_ACC_Z0;
    H[2][2] = -2 * q2 * IMU_ACC_Z0;
    H[3][2] = -2 * q2 * IMU_MAG_B0[0][0] + 2 * q1 * IMU_MAG_B0[1][0] - 2 * q0 * IMU_MAG_B0[2][0];
    H[4][2] =  2 * q1 * IMU_MAG_B0[0][0] + 2 * q2 * IMU_MAG_B0[1][0] + 2 * q3 * IMU_MAG_B0[2][0];
    H[5][2] =  2 * q0 * IMU_MAG_B0[0][0] + 2 * q3 * IMU_MAG_B0[1][0] - 2 * q2 * IMU_MAG_B0[2][0];
    
    H[0][3] = +2 * q1 * IMU_ACC_Z0;
    H[1][3] = +2 * q2 * IMU_ACC_Z0;
    H[2][3] = +2 * q3 * IMU_ACC_Z0;
    H[3][3] = -2 * q3 * IMU_MAG_B0[0][0] + 2 * q0 * IMU_MAG_B0[1][0] + 2 * q1 * IMU_MAG_B0[2][0];
    H[4][3] = -2 * q0 * IMU_MAG_B0[0][0] - 2 * q3 * IMU_MAG_B0[1][0] + 2 * q2 * IMU_MAG_B0[2][0];
    H[5][3] =  2 * q1 * IMU_MAG_B0[0][0] + 2 * q2 * IMU_MAG_B0[1][0] + 2 * q3 * IMU_MAG_B0[2][0];
    
    return true;
}