#include "kalman.h"
#include <stdlib.h>
#include "euler.h"

void SetDiagonal(float *matr, int m, int n, float value) {
  int k, i;
  k = (m<n ? m : n);
  for (i=0; i<m*n; i++) matr[i] = 0.0;
  for (i=0; i<k; i++) matr[i*n + i] = value;
}

void Transpose (double R[9], double S[9]) {
  int i, j;
  for (i=0; i<3; i++)
   for (j=0; j<3; j++) 
    S[3*i+j] = R[3*j+i];
}

void Mult (double R[9], double S[9], double T[9]) {
  int i, j;
  for (i=0; i<3; i++)
   for (j=0; j<3; j++) 
    T[3*i+j] = R[3*i+0]*S[0*3+j] + R[3*i+1]*S[1*3+j]+ R[3*i+2]*S[2*3+j];
}

KalmanStruct * init_kf_points(int npts, double *xf, double *yf, 
                             float sigmaR) {
  float *Phi, *H, *Q, *R, *x, *P;
  KalmanStruct *kf;
  int i;

  Phi = (float *) malloc(16*npts*npts*sizeof(float));
  H = (float *) malloc(8*npts*npts*sizeof(float));
  Q = (float *) malloc(16*npts*npts*sizeof(float));
  R = (float *) malloc(4*npts*npts*sizeof(float));
  P = (float *) malloc(16*npts*npts*sizeof(float));
  x = (float *) malloc(4*npts*sizeof(float));

  SetDiagonal(Phi, 4*npts, 4*npts, 1.0);
  for (i=0; i<2*npts; i++) Phi[4*npts*i + 2*npts + i] = 1;
  SetDiagonal(H, 2*npts, 4*npts, 1.0);
  SetDiagonal(Q, 4*npts, 4*npts, 0.0);
  for (i=2*npts; i<4*npts; i++) {Q[i*4*npts + i] = 1.0;}
  SetDiagonal(R, 2*npts, 2*npts, sigmaR);
  SetDiagonal(P, 4*npts, 4*npts, 1e5*sigmaR);
  for (i=0; i<npts; i++) {x[i] = xf[i]; x[npts+i] = yf[i]; 
                          x[2*npts+ 2*i] = 0.0;  x[2*npts + 2*i + 1] = 0.0;}
  kf = CreateKalman(4*npts, 2*npts);
  kf->Phi = Phi;
  kf->Q = Q;
  kf->R = R;
  kf->H = H;
  kf->x = x;
  kf->P = P;
  return (kf);

}

int step_kf_points(int npts, double *xf, double *yf, float sigmaQ, int dt, KalmanStruct *kf)   {
  float *z;
  int i;

  z = (float *) malloc(2*npts*sizeof(float));
  for (i=0; i<npts; i++) {z[i] = xf[i]; z[npts+i] = yf[i];}
  for (i=2*npts; i<4*npts; i++) kf->Q[i*4*npts + i] = sigmaQ *dt;
  for (i=0; i<2*npts; i++) kf->Phi[4*npts*i + 2*npts + i] = dt;
  KalmanPredict(kf);
  KalmanCorrect(kf, z);
  free(z);
  for (i=0; i<2*npts; i++) kf->Phi[4*npts*i + 2*npts + i] = dt;
  for (i=0; i<npts; i++) {xf[i] = kf->x[i]; yf[i]= kf->x[npts+i] ;}
  if (kf->chi_square < 50 ) {
     for (i=0; i<npts; i++) {xf[i] = kf->x[i]; yf[i]= kf->x[npts+i] ;}
     return 1;
  }
  else {
//    for (i=0; i<4*npts; i++) kf->x[i] = kf->PriorState[i];
//    for (i=0; i<16*npts*npts; i++) kf->P[i] = kf->PriorErrorCovariance[i];
    for (i=0; i<npts; i++) {xf[i] = kf->x[i]; yf[i]= kf->x[npts+i] ;}
    return 0;
  }
}

KalmanStruct * init_kf_camera(double sigmaR, double sigmaT, double *r, double *T) {
  float *Phi, *H, *Q, *R, *x, *P;
  KalmanStruct *kf;
  int i; 
  double w[3];

  Phi = (float *) malloc(12*12*sizeof(float));
  H = (float *) malloc(12*6*sizeof(float));
  Q = (float *) malloc(12*12*sizeof(float));
  R = (float *) malloc(6*6*sizeof(float));
  P = (float *) malloc(12*12*sizeof(float));
  x = (float *) malloc(12*sizeof(float));

  SetDiagonal(Phi, 12, 12, 1.0);
  for (i=0; i<6; i++) Phi[12*i + 6 + i] = 1;
  SetDiagonal(H, 6, 12, 1.0);
  SetDiagonal(Q, 12, 12, 0.0);
  for (i=6; i<12; i++) {Q[12*i + i] = 1;}
  SetDiagonal(R, 6, 6, sigmaT);
  for (i=0; i<3; i++) R[6*i+i] = sigmaR;
  SetDiagonal(P, 12, 12, 100000);
  for (i=0; i<6; i++) P[12*i+i] = 0.00001;
  rotation_to_angles(r, w);
  x[0] = w[0]; x[1] = w[1]; x[2] = w[2];
  x[3] = T[0]; x[4] = T[1]; x[5] = T[2];
  for (i=6; i<12; i++) x[i] = 0.0;
  kf = CreateKalman(12, 6);
  kf->Phi = Phi;
  kf->Q = Q;
  kf->R = R;
  kf->H = H;
  kf->x = x;
  kf->P = P;
  return (kf);

}

int step_kf_camera(double sigmaR, double sigmaT, double *r, double *T,  double dt, KalmanStruct *kf) {

  float z[6];
  int i;
  double w[3],  w0[3],  R0[9],  Rincr[9], R0inv[9];
  
  w0[0] = kf->x[0]; w0[1] = kf->x[1]; w0[2] = kf->x[2];
  angles_to_rotation(w0, R0);
  Transpose(R0, R0inv);
  Mult (R0inv, r, Rincr);
  rotation_to_angles(Rincr, w);
  kf->x[0] = kf->x[1] = kf->x[2] = 0.0;
  z[0] = w[0] ;
  z[1] = w[1] ;
  z[2] = w[2] ;
  z[3] = T[0]; z[4] = T[1]; z[5] = T[2]; 
  for (i=6; i<9; i++) kf->Q[i*12 + i] = sigmaR *dt;
  for (i=9; i<12; i++) kf->Q[i*12 + i] = sigmaT *dt;
  for (i=0; i<6; i++) kf->Phi[12*i + 6 + i] = dt;
  KalmanPredict(kf);
  KalmanCorrect(kf, z);
  w[0] = kf->x[0];  
  w[1] = kf->x[1];  
  w[2] = kf->x[2];
  angles_to_rotation(w, Rincr);
  Mult(R0, Rincr, r);
  rotation_to_angles(r, w);
  kf->x[0] = w[0]; kf->x[1] = w[1]; kf->x[2] = w[2]; 
  T[0] = kf->x[3]; T[1] = kf->x[4]; T[2] = kf->x[5];
  return (1);

}

