#include <BasicLinearAlgebra.h>

class Kalman2DFilter {
  public:
    Kalman2DFilter(float dt = 4.f, float accelUncertainty = 10.f, float barometerUncertainty = 100.f) {
        dt /= 1000.f;
        F = { 1, dt,
              0, 1 };
        G = { 0.5f * dt * dt,
              dt };
        H = { 1, 0 };
        I = { 1, 0,
              0, 1 };
        Q = G * ~G * accelUncertainty * accelUncertainty;
        R = { barometerUncertainty * barometerUncertainty };
        P = { 0, 0,
              0, 0 };
        S = { 0,
              0 };
    }

    void filter(const float &AccZInertial,
                const float &AltitudeBarometer,
                float &AltitudeKalman,
                float &VelocityVerticalKalman) {
        Acc = { AccZInertial };
        S = F * S + G * Acc;
        P = F * P * ~F + Q;
        L = H * P * ~H + R;
        K = P * ~H * Inverse(L);
        M = { AltitudeBarometer };
        S = S + K * (M - H * S);
        AltitudeKalman = S(0, 0);
        VelocityVerticalKalman = S(1, 0);
        P = (I - K * H) * P;
    }

  private:
    BLA::Matrix<2, 2> F;
    BLA::Matrix<2, 1> G;
    BLA::Matrix<2, 2> P;
    BLA::Matrix<2, 2> Q;
    BLA::Matrix<2, 1> S;
    BLA::Matrix<1, 2> H;
    BLA::Matrix<2, 2> I;
    BLA::Matrix<1, 1> Acc;
    BLA::Matrix<2, 1> K;
    BLA::Matrix<1, 1> R;
    BLA::Matrix<1, 1> L;
    BLA::Matrix<1, 1> M;
};