#include "FtmKalman.h"

#include <iostream>
#include <eigen3/Eigen/Eigen>

constexpr auto square(float x) { return x * x; };


float Kalman::predictAndUpdate(const Timestamp timestamp, const float measurment)
{
    const auto I = Eigen::Matrix2f::Identity();

    {
    // hack
        processNoiseDistance = 1.2f; //1.2f;
        processNoiseVelocity = 1.5f; //1.5;
        R = 300;// 200;
    }

    Eigen::Map<Eigen::Matrix<float, 2, 1>> x(this->x);
    Eigen::Map<Eigen::Matrix<float, 2, 2>> P(this->P);

    // init kalman filter
    if (std::isnan(lastTimestamp))
    {
        P << 10, 0,
            0, 10; // Initial Uncertainty   

        x << measurment,
            0;
    }

    const float dt = std::isnan(lastTimestamp) ? 1 : timestamp.sec() - lastTimestamp;
    lastTimestamp = timestamp.sec();

    Eigen::Matrix<float, 1, 2> H; // Measurement function
    H << 1, 0;
    
    Eigen::Matrix2f A;   // Transition Matrix
    A << 1, dt,
        0, 1;

    Eigen::Matrix2f Q;  // Process Noise Covariance
    Q << square(processNoiseDistance), 0,
        0, square(processNoiseVelocity);

    // Prediction
    x = A * x;                    // Prädizierter Zustand aus Bisherigem und System
    P = A * P*A.transpose() + Q;    // Prädizieren der Kovarianz

    // Correction
    float Z = measurment;
    auto y = Z - (H*x);                // Innovation aus Messwertdifferenz
    auto S = (H*P*H.transpose() + R);    // Innovationskovarianz
    auto K = P * H.transpose()* (1 / S); //Filter-Matrix (Kalman-Gain)

    x = x + (K*y);     // aktualisieren des Systemzustands
    P = (I - (K*H))*P; // aktualisieren der Kovarianz

    return x(0);
}

KalmanPrediction Kalman::predict(const Timestamp timestamp)
{
    if (std::isnan(lastTimestamp))
    {
        // Kalman has no data => nothing to predict
        return KalmanPrediction{ NAN, NAN };
    }
    else 
    {
        Eigen::Map<Eigen::Matrix<float, 2, 1>> x(this->x);
        Eigen::Map<Eigen::Matrix<float, 2, 2>> P(this->P);

        const float dt = timestamp.sec() - lastTimestamp;
        lastTimestamp = timestamp.sec();

        Eigen::Matrix2f A;   // Transition Matrix
        A << 1, dt,
             0, 1;


        Eigen::Matrix2f Q;  // Process Noise Covariance
        Q << square(processNoiseDistance), 0,
             0, square(processNoiseVelocity);


        x = A * x;
        P = A * P*A.transpose() + Q;


        return KalmanPrediction{ x(0), x(1) };
    }
}