//#line 2 "kf.c"
#include <cblas.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#include "kf.h"

// Follow along in Welch & Bishops "An Introduction to the Kalman Filter"

void printMatrixKF(double* M, int row, int col) 
{
    int i, j;
    for (i=0; i<row; ++i) {
        for (j=0; j<col; ++j)  fprintf(stderr,"%+0.4f\t", M[i*col+j]);
        fprintf(stderr,"\n");
    }
}

kf_struct initKF(int n, int m, int l) 
{
    int i;
    kf_struct kf;
    kf.n = n;
    kf.m = m;
    kf.l = l;
    kf.x = (double*) calloc(n, sizeof(double));
    kf.z = (double*) calloc(m, sizeof(double));
    kf.u = (double*) calloc(l, sizeof(double));
    kf.P = (double*) calloc(n*n, sizeof(double));
    kf.F = (double*) calloc(n*n, sizeof(double));
    kf.Q = (double*) calloc(n*n, sizeof(double));
    kf.G = (double*) calloc(n*l, sizeof(double));
    kf.H = (double*) calloc(m*n, sizeof(double));
    kf.R = (double*) calloc(m*m, sizeof(double));
    for (i=0;i<n;i++) {
        kf.P[i*n+i] = 1;
        kf.F[i*n+i] = 1;
        kf.Q[i*n+i] = 1;
    }
    kf.G[0] = 1;
    kf.H[0] = 1;
    kf.R[0] = 1;
    return kf;
}

void deleteKF(kf_struct kf) 
{
    free(kf.x);
    free(kf.z);
    free(kf.u);
    free(kf.P);
    free(kf.F);
    free(kf.Q);
    free(kf.G);
    free(kf.H);
    free(kf.R);
}

void printKF(kf_struct kf) 
{
    // fprintf(stderr, "n = %d\n", kf.n);
    fprintf(stderr, "x = \n");
    printMatrixKF(kf.x, kf.n, 1);
    fprintf(stderr, "u = \n");
    printMatrixKF(kf.u, kf.m, 1);
    fprintf(stderr, "z = \n");
    printMatrixKF(kf.z, kf.m, 1);
    fprintf(stderr, "F = \n");
    printMatrixKF(kf.F, kf.n, kf.n);
//    fprintf(stderr, "G = \n");
//    printMatrixKF(kf.G, kf.n, kf.m);
    fprintf(stderr, "Q = \n");
    printMatrixKF(kf.Q, kf.n, kf.n);
    fprintf(stderr, "P = \n");
    printMatrixKF(kf.P, kf.n, kf.n);
//    fprintf(stderr, "H = \n");
//    printMatrixKF(kf.H, kf.n, 1);
    fprintf(stderr, "R = \n");
    printMatrixKF(kf.R, kf.m, kf.m);
}

void predictKF(kf_struct kf)
{
    int i;
    // Predict state: x = Fx + Gu
    double Fx[kf.n];
    cblas_dgemv(CblasRowMajor, CblasNoTrans, kf.n, kf.n, 
                1, kf.F, kf.n, kf.x, 1, 0, Fx, 1);
    double Gu[kf.n];
    cblas_dgemv(CblasRowMajor, CblasNoTrans, kf.n, kf.l, 
                1, kf.G, kf.l, kf.u, 1, 0, Gu, 1);
    for (i = 0; i<kf.n; ++i) {
        kf.x[i] = Fx[i] + Gu[i];
    }
    // Predict estimate covariance:  P = FPF' + Q
    double FP[kf.n*kf.n];
    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                kf.n, kf.n, kf.n, 1, kf.F, kf.n, kf.P, kf.n, 0, FP, kf.n);
    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasTrans,
                kf.n, kf.n, kf.n, 1, FP, kf.n, kf.F, kf.n, 0, kf.P, kf.n);
    for (i = 0; i<kf.n*kf.n; ++i) {
        kf.P[i] += kf.Q[i];
    }
    //fprintf(stderr,"P=\n");
    //printMatrixKF(kf.P, kf.n, kf.n);
}

void correctKF(kf_struct kf)
{
    int i;
    // Innovation residual:  y = z - Hx
    double Hx[kf.m];
    cblas_dgemv(CblasRowMajor, CblasNoTrans, kf.m, kf.n,
                1, kf.H, kf.n, kf.x, 1, 0, Hx, 1);
    double y[kf.m];
    for (i = 0; i<kf.m; ++i) {
        y[i] = kf.z[i] - Hx[i];
    }

    // Innovation covariance: S = HPH' + R
    double HP[kf.m * kf.n];
    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                kf.m, kf.n, kf.n, 1, kf.H, kf.n, kf.P, kf.n, 0, HP, kf.n);
    double HPH[kf.m * kf.m];
    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasTrans,
                kf.m, kf.m, kf.n, 1, HP, kf.n, kf.H, kf.n, 0, HPH, kf.m);
    double S[kf.m * kf.m];
    for (i = 0; i<kf.m * kf.m; ++i) {
        S[i] = HPH[i] + kf.R[i];
    }

    //fprintf(stderr,"S= %g\n", S[0]);

    // Kalman gain: K = PH'S^-1
    double K[kf.n * kf.m];
    // Hack for now!
    double PH[kf.n * kf.m];
    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasTrans,
                kf.n, kf.m, kf.n, 1, kf.P, kf.n, kf.H, kf.n, 0, PH, kf.m);

    for (i = 0; i<kf.n*kf.m; ++i) {
        K[i] = PH[i] / S[0];
    }

    //fprintf(stderr,"K=\n");
    //printMatrixKF(K,kf.n,kf.m);

    // Updated state estimate: x = x + Ky
    double Ky[kf.n];
    cblas_dgemv(CblasRowMajor, CblasNoTrans, kf.n, kf.m,
                1, K, kf.m, y, 1, 0, Ky, 1);
    for (i = 0; i<kf.n; ++i) {
        kf.x[i] += Ky[i];
    }

    // Update estimate covariance: P = P - KHP
    double KH[kf.n * kf.n];
    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                kf.n, kf.n, kf.m, 1, K, kf.m, kf.H, kf.n, 0, KH, kf.n);
    double KHP[kf.n*kf.n];
    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                kf.n, kf.n, kf.n, 1, KH, kf.n, kf.P, kf.n, 0, KHP, kf.n);
    for (i = 0; i<kf.n*kf.n; ++i) {
        kf.P[i] -= KHP[i];
    }
    //fprintf(stderr,"P=\n");
    //printMatrixKF(kf.P, kf.n, kf.n);
}