#include "OrientationFusion.h"
#include "time/SystemTime.h"

using namespace Sensors;

#define matrix_set(matrix, row, column, value) \
    matrix->data[row][column] = value

#define matrix_set_symmetric(matrix, row, column, value) \
    matrix->data[row][column] = value; \
    matrix->data[column][row] = value

#define matrix_get(matrix, row, column) \
    matrix->data[row][column]


#ifndef FIXMATRIX_MAX_SIZE
#error FIXMATRIX_MAX_SIZE must be defined and greater or equal to the number of states, inputs and measurements.
#endif

OrientationFusion::OrientationFusion()
{
#if (FIXMATRIX_MAX_SIZE < States) || (FIXMATRIX_MAX_SIZE < Inputs) || (FIXMATRIX_MAX_SIZE < Measurements)
#error FIXMATRIX_MAX_SIZE must be greater or equal to the number of states, inputs and measurements.
#endif

    /************************************************************************/
    /* initialize the filter structures                                     */
    /************************************************************************/
    kalman_filter_initialize_uc(&kf, States);
    kalman_observation_initialize(&kfm, States, Measurements);

    /************************************************************************/
    /* set initial state                                                    */
    /************************************************************************/
    mf16 *x = kalman_get_state_vector_uc(&kf);

    x->data[0][0] = 0; // theta
    x->data[1][0] = 0; // omega
    x->data[2][0] = 0; // alpha

    /************************************************************************/
    /* set covariance                                                       */
    /************************************************************************/
    mf16 *P = kalman_get_system_covariance_uc(&kf);
    matrix_set_symmetric(P, 0, 0, F16(.029792 * .029792));   // var(v)
    matrix_set_symmetric(P, 0, 1, 0);   // cov(v,a)
    matrix_set_symmetric(P, 0, 2, 0);   // cov(v,j)

    matrix_set_symmetric(P, 1, 1, F16(.00857434 * .00857434));   // var(a)
    matrix_set_symmetric(P, 1, 2, 0);   // cov(a,j)

    matrix_set_symmetric(P, 2, 2, F16(.0085743 * .0085743));   // var(j)

    /************************************************************************/
    /* system process noise                                                 */
    /************************************************************************/
    mf16 *Q = kalman_get_system_process_noise_uc(&kf);
    matrix_set_symmetric(Q, 0, 0, F16(.029792 * .029792));   // var(v)
    matrix_set_symmetric(Q, 0, 1, 0);   // cov(v,a)
    matrix_set_symmetric(Q, 0, 2, 0);   // cov(v,j)

    matrix_set_symmetric(Q, 1, 1, F16(.129792 * .129792));   // var(a)
    matrix_set_symmetric(Q, 1, 2, 0);   // cov(a,j)

    matrix_set_symmetric(Q, 2, 2, F16(.0085743 * .0085743));   // var(j)

    /************************************************************************/
    /* set measurement transformation                                       */
    /************************************************************************/
    mf16 *H = kalman_get_observation_transformation(&kfm);

    matrix_set(H, 0, 0, fix16_one); // z = 1*theta
    matrix_set(H, 0, 1, 0);         //   + 0*omega
    matrix_set(H, 0, 2, 0);         //   + 0*alpha
    matrix_set(H, 1, 0, 0);         // z = 0*theta
    matrix_set(H, 1, 1, fix16_one); //   + 1*omega
    matrix_set(H, 1, 2, 0);         //   + 0*alpha

    /************************************************************************/
    /* set process noise                                                    */
    /************************************************************************/
    mf16 *R = kalman_get_observation_process_noise(&kfm);

    matrix_set(R, 0, 0, F16(3*0.129792*0.129792));     // var(theta)
    matrix_set(R, 1, 1, F16(0.0085743*0.0085743));     // var(omega)
}

void OrientationFusion::SetStateTransition(milliseconds_t time)
{
    mf16 *A = kalman_get_state_transition_uc(&kf);

    // set time constant
    const fix16_t T = fix16_div(fix16_from_int(time.value), fix16_from_int(1000));
    const fix16_t Tsquare = fix16_sq(T);

    // helper
    const fix16_t fix16_half = fix16_from_float(0.5);

    // transition of x to v
    matrix_set(A, 0, 0, fix16_one);   // 1
    matrix_set(A, 0, 1, T);   // T
    matrix_set(A, 0, 2, fix16_mul(fix16_half, Tsquare)); // 0.5 * T^2

    // transition of x to a
    matrix_set(A, 1, 0, 0);   // 0
    matrix_set(A, 1, 1, fix16_one);   // 1
    matrix_set(A, 1, 2, T);   // T

    // transition of x to j
    matrix_set(A, 2, 0, 0);   // 0
    matrix_set(A, 2, 1, 0);   // 0
    matrix_set(A, 2, 2, fix16_one);   // 1
}

milliseconds_t OrientationFusion::Predict(Fix16 u_omega)
{
    const auto now = SystemTime::get()->NowMs();
    const auto difference = (now - _previousRun);
    _previousRun = now;

    SetStateTransition(difference);
    kalman_predict_uc(&kf);

    return difference;
}


void OrientationFusion::Update(milliseconds_t dt, Fix16 z_heading, Fix16 z_odometry)
{
    /* update */
    mf16 *z = kalman_get_observation_vector(&kfm);
    matrix_set(z, 0, 0, z_heading);
    matrix_set(z, 1, 0, z_odometry);

    /* correct */
    kalman_correct_uc(&kf, &kfm);

    mf16 *P = kalman_get_system_covariance_uc(&kf);
    matrix_set_symmetric(P, 0, 1, 0);
    matrix_set_symmetric(P, 0, 2, 0);
    matrix_set_symmetric(P, 1, 2, 0);

    if (kf.x.data[1][0] == 0xFFFFFFFF80000000 || kf.x.data[1][0] == 0x80000000)
    {
        kf.x.data[0][0] = 0;
        kf.P.data[0][0] = Fix16::OneThousand;

        kf.x.data[1][0] = 0;
        kf.P.data[1][1] = Fix16::OneThousand;

        kf.x.data[2][0] = 0;
        kf.P.data[2][2] = Fix16::OneThousand;
    }
}