#include "kalman.h"
#include <stddef.h>

// Perform the Kalman filter prediction step.
// X: State column vector
// S: Square root of state covariance
// F: State transition matrix or NULL for identity
// Q: Process noise covariance matrix or NULL for no noise
void kalman_prediction_step(mf16 *X, mf16 *S, const mf16 *F, const mf16 *Q)
{
    // tmp = (F * S)'
    mf16 tmp = {};
    if (F != NULL)
    {
        mf16_mul(&tmp, F, S);
        mf16_transpose(&tmp, &tmp);
    }
    else
    {
        mf16_transpose(&tmp, S);
    }
    
    // Copy Q' to lower part of tmp
    int n = tmp.rows;
    for (int row = 0; row < n; row++)
    {
        for (int col = 0; col < n; col++)
        {
            if (Q != NULL)
            {
                tmp.data[n + col][row] = Q->data[row][col];
            }
            else
            {
                tmp.data[n + col][row] = 0;
            }
        }
    }
    tmp.rows *= 2;
    
    // QR-decompose tmp and R is the new S'
    mf16_qr_decomposition(&tmp, S, &tmp, 0);
    mf16_transpose(S, S);
    
    // X = F * X
    if (F != NULL)
        mf16_mul(X, F, X);
}

// Perform the Kalman filter update step for one measurement variable.
// X: State column vector
// S: Square root of state covariance
// H: Measurement model (derivative) matrix
// R: Measurement standard deviation (noise)
// Y: Residual value of the measurement (measurement - prediction)
void kalman_potter_update(mf16 *X, mf16 *S, const mf16 *H, fix16_t R, fix16_t Y)
{
    R = fix16_mul(R, R);
    
    fix16_t C;
    {
        // C = H S S' H' + R
        mf16 HS;
        mf16_mul(&HS, H, S); // Calculate H S
        mf16_mul_bt(&HS, &HS, &HS); // Calculate H S (H S)'
        C = HS.data[0][0] + R; // Result is scalar.
        X->errors |= HS.errors;
    }
    
    mf16 K;
    {
        // K = S S' H' / C
        mf16_mul(&K, H, S); // Calculate H S
        mf16_mul_bt(&K, S, &K); // Calculate S S' H'
        mf16_div_s(&K, &K, C);
    }
    
    {
        // X_new = X + K Y
        mf16 tmp;
        mf16_mul_s(&tmp, &K, Y);
        mf16_add(X, X, &tmp);
    }
    
    {
        // S_new = (I - K H / d) S = S - K H S / d
        // d = 1 + R / C
        fix16_t d = fix16_one + fix16_div(R, C);
        mf16 tmp;
        mf16_mul(&tmp, &K, H); // Compute K H
        mf16_mul(&tmp, &tmp, S); // Compute K H S
        mf16_div_s(&tmp, &tmp, d); // Compute K H S / d
        mf16_sub(S, S, &tmp);
    }
}

// Perform the Kalman filter update step.
// X: State column vector
// S: Square root of state covariance
// H: Measurement model (derivative) matrix
// R: Measurement noise covariance matrix
// Y: Residual values of the measurements, column vector
void kalman_update(mf16 *X, mf16 *S, const mf16 *H, const mf16 *R, const mf16 *Y)
{
    // C = H S S' H' + R
    mf16 C;
    mf16_mul(&C, H, S); // Calculate H S
    mf16_mul_bt(&C, &C, &C); // Calculate H S (H S)'
    mf16_add(&C, &C, R); // Add R
    
    // K = S S' H' inv(C)
    // K' = inv(C') H S S'
    mf16 K, Cq, Cr;
    mf16_transpose(&C, &C); // Transpose C to C'
    mf16_mul(&K, H, S); // Calculate H S
    mf16_mul_bt(&K, &K, S); // Calculate H S S'
    mf16_qr_decomposition(&Cq, &Cr, &C, 0); // Decompose C for solving
    mf16_solve(&K, &Cq, &Cr, &K); // Solve K'
    mf16_transpose(&K, &K); // Transpose K' to K
    
    // X = X + K y
    mf16 tmp;
    mf16_mul(&tmp, &K, Y);
    mf16_add(X, X, &tmp);
    
    // P+ = (I - K H) S S' (I - K H)' + K R K'
    mf16 tmp2;
    tmp.rows = tmp.columns = K.rows;
    mf16_fill_diagonal(&tmp, fix16_from_int(1)); // tmp = I
    mf16_mul(&tmp2, &K, H); // tmp2 = K H
    mf16_sub(&tmp, &tmp, &tmp2); // tmp = I - K H
    mf16_mul(S, &tmp, S); // Calculate (I - K H) S
    mf16_mul_bt(S, S, S); // Calculate (I - K H) S ((I - K H) S)'
    mf16_mul(&tmp, &K, R); // Calculate K R
    mf16_mul_bt(&tmp, &tmp, &K); // Calculate K R K'
    mf16_add(S, S, &tmp); // Now we have P+
    mf16_cholesky(S, S); // Decompose P+ to S+
}