#include "libscl.h"
#include "svmod.h"
using namespace scl; 
using namespace std;

int main(int argc, char** argp, char** envp)
{
  INTEGER n = 200;
  INTEGER N = 10000;
  INT_32BIT seed = 780695;
  
  svmod m;

  sample s = m.draw_sample(n,seed);

  INTEGER rows_X  = s.x0.nrow();
  INTEGER rows_Y  = s.y0.nrow();

  if ((rows_Y%2 != 0) || (s.Y(1,1) != s.Y(rows_Y/2+1,2))) {
    error("Error, assumption about state vector wrong");
    // The assumption is that the first row(s) of y are time t and second t-1
  }

  vector<realmat> buffer1(N);
  vector<realmat> buffer2(N);
  vector<realmat> filter(N);

  vector<realmat>* smooth = &buffer1;
  vector<realmat>* draws = &buffer2;

  REAL weights[N];
  REAL log_likelihood = 0.0;

  // Initialization

  realmat Y(rows_Y,n+1);

  for (INTEGER i=1; i<=rows_Y/2; ++i) Y(rows_Y/2+i,1) = 0.0;
  for (INTEGER i=1; i<=rows_Y/2; ++i) Y(i,1) = s.Y(rows_Y/2+i,1);
  for (INTEGER t=2; t<=n+1; t++) {
    for (INTEGER i=1; i<=rows_Y; ++i) Y(i,t) = s.Y(i,t-1);
  }

  realmat X(rows_X,n+1);

  realmat x0;
  realmat y0;
  REAL sum = 0.0;

  for (INTEGER i=0; i<N; ++i) {
    m.draw_x0_y0(x0,y0,seed);
    for (INTEGER k=1; k<=(rows_Y/2); ++k) y0[k] = Y(k,1);
    sum += m.prob_yt(y0,x0);
    for (INTEGER k=1; k<=rows_X; ++k) {
      X(k,1) = x0[k];
    }
    (*smooth)[i] = X;
    (*draws)[i] = X;
    filter[i] = X;
  }
  log_likelihood += log(sum/REAL(N));

  realmat xlag(rows_X,1);
  realmat xt(rows_X,1);
  realmat yt(rows_Y,1);
  for (INTEGER t=2; t<=n+1; ++t) {

    vector<realmat>* save = draws;
    draws = smooth;
    smooth = save;

    // Importance sampling step

    sum = 0.0;
    for (INTEGER i=0; i<N; ++i) {
      for (INTEGER k=1; k<=rows_X; ++k) {
        xlag[k] = (*draws)[i](k,t-1);
      }
      xt = m.draw_xt(xlag,seed);
      for (INTEGER k=1; k<=rows_Y; ++k) {
        yt[k] = Y(k,t);
      }
      for (INTEGER k=1; k<=rows_X; ++k) {
        (*draws)[i](k,t) = xt[k];
      }
      sum += weights[i] = m.prob_yt(yt,xt);
    }
    log_likelihood += log(sum/REAL(N));
  
    weights[0] /= sum;
    for (INTEGER i=1; i<(N-1); ++i) {
      weights[i] /= sum;
      weights[i] += weights[i-1];
    }
    weights[N-1] = 1.0;

    // Selection step

    for (INTEGER i=0; i<N; ++i) {
      REAL u = ran(seed);
      INTEGER j = 0;
      while(weights[j] <= u) ++j;
      (*smooth)[i] = (*draws)[j];
      for (INTEGER k=1; k<=rows_X; ++k) {
        filter[i](k,t) = (*draws)[j](k,t);
      }
    }
  }  
    
  realmat mean;
  realmat sdev;
  ofstream fout;

  mean.resize(rows_X,n+1,0.0);
  for (INTEGER i=0; i<N; ++i) {
    mean += (*smooth)[i];
  }
  mean = mean/N;

  sdev.resize(rows_X,n+1,0.0);
  for (INTEGER i=0; i<N; ++i) {
    realmat z = (*smooth)[i] - mean;
    for (INTEGER t=1; t<=n+1; ++t) {
      for (INTEGER k=1; k<=rows_X; ++k) {
        sdev(k,t) += z(k,t)*z(k,t);
      }
    }
  }
  for (INTEGER t=1; t<=n+1; ++t) {
    for (INTEGER k=1; k<=rows_X; ++k) {
      sdev(k,t) = sqrt(sdev(k,t)/REAL(N-1));
    }
  }
  
  fout.open("smooth.csv");
  if (!fout) error("Error, particle, cannot open fout");

  for (INTEGER k=1; k<=rows_X; ++k) fout << "mean" << k << ',';
  for (INTEGER k=1; k<=rows_X; ++k) fout << "sdev" << k << ',';
  for (INTEGER k=1; k<=rows_X; ++k) fout << "x" << k << ',';
  for (INTEGER k=1; k<=rows_Y-1; ++k) fout << "y" << k << ',';
  fout << "y" << rows_Y << '\n';
  
  for (INTEGER k=1; k<=rows_X; ++k) fout << mean(k,1) << ',';
  for (INTEGER k=1; k<=rows_X; ++k) fout << sdev(k,1) << ',';
  for (INTEGER k=1; k<=rows_X; ++k) fout << s.x0[k] << ',';
  for (INTEGER k=1; k<=rows_Y-1; ++k) fout << s.x0[k] << ',';
  fout << s.x0[rows_Y] << '\n';

  for (INTEGER t=2; t<=n+1; ++t) {
    for (INTEGER k=1; k<=rows_X; ++k) fout << mean(k,t) << ',';
    for (INTEGER k=1; k<=rows_X; ++k) fout << sdev(k,t) << ',';
    for (INTEGER k=1; k<=rows_X; ++k) fout << s.X(k,t-1)  << ',';
  for (INTEGER k=1; k<=rows_Y-1; ++k) fout << s.Y(k,t-1) << ',';
  fout << s.Y(rows_Y-1,t-1) << '\n';
  }
  
  fout.clear(); fout.close();

  mean.resize(rows_X,n+1,0.0);
  for (INTEGER i=0; i<N; ++i) {
    mean += filter[i];
  }
  mean = mean/N;

  sdev.resize(rows_X,n+1,0.0);
  for (INTEGER i=0; i<N; ++i) {
    realmat z = filter[i] - mean;
    for (INTEGER t=1; t<=n+1; ++t) {
      for (INTEGER k=1; k<=rows_X; ++k) {
        sdev(k,t) += z(k,t)*z(k,t);
      }
    }
  }
  for (INTEGER t=1; t<=n+1; ++t) {
    for (INTEGER k=1; k<=rows_X; ++k) {
      sdev(k,t) = sqrt(sdev(k,t)/REAL(N-1));
    }
  }
  
  fout.open("filter.csv");
  if (!fout) error("Error, particle, cannot open fout");

  for (INTEGER k=1; k<=rows_X; ++k) fout << "mean" << k << ',';
  for (INTEGER k=1; k<=rows_X; ++k) fout << "sdev" << k << ',';
  for (INTEGER k=1; k<=rows_X; ++k) fout << "x" << k << ',';
  for (INTEGER k=1; k<=rows_Y-1; ++k) fout << "y" << k << ',';
  fout << "y" << rows_Y << '\n';
  
  for (INTEGER k=1; k<=rows_X; ++k) fout << mean(k,1) << ',';
  for (INTEGER k=1; k<=rows_X; ++k) fout << sdev(k,1) << ',';
  for (INTEGER k=1; k<=rows_X; ++k) fout << s.x0[k] << ',';
  for (INTEGER k=1; k<=rows_Y-1; ++k) fout << s.x0[k] << ',';
  fout << s.x0[rows_Y] << '\n';

  for (INTEGER t=2; t<=n+1; ++t) {
    for (INTEGER k=1; k<=rows_X; ++k) fout << mean(k,t) << ',';
    for (INTEGER k=1; k<=rows_X; ++k) fout << sdev(k,t) << ',';
    for (INTEGER k=1; k<=rows_X; ++k) fout << s.X(k,t-1)  << ',';
  for (INTEGER k=1; k<=rows_Y-1; ++k) fout << s.Y(k,t-1) << ',';
  fout << s.Y(rows_Y-1,t-1) << '\n';
  }

  fout.clear(); fout.close();

  cout << "The log likelihood is " << log_likelihood << '\n';

  return 0;
}