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

#define DEBUG_CHECKS

struct sample {
  scl::realmat X;
  scl::realmat Y;
  scl::realmat X0_Y0;
  scl::realmat Mean_X0_Y0;
  scl::realmat Var_X0_Y0;
};

class pfmod : public scl::pfmod_base {
private:
  REAL a0,a1;
  REAL b01,b02;
  REAL b11,b21,b12,b22;
  REAL r11,r21,r22,r31,r32,r33;
  INTEGER xdim, xlag;
  INTEGER ydim, ylag;
  scl::denval psave;
  scl::realmat Mean_X0_Y0;
  scl::realmat Var_X0_Y0;
public:
  pfmod() 
  : a0(0), a1(0.1), 
    b01(0), b02(0), 
    b11(0.01), b21(0.0), b12(0.0), b22(0.9), 
    r11(0.25), r21(0.0), r22(0.25), r31(0.01), r32(0.0), r33(0.5), 
    xdim(2), xlag(1), ydim(1), ylag(1), psave() 
    { scl::realmat theta = get_parms(); set_parms(theta); }
  void set_parms(const scl::realmat& theta); 
  scl::realmat get_parms();
  void q_draw_x0_y0(scl::particles* pf, INT_32BIT& seed);
  scl::denval q_prob_x0_y0(scl::particles* pf);
  scl::denval p_prob_x0_y0(scl::particles* pf);
  void q_draw_xt(scl::particles* pf, INT_32BIT& seed);
  scl::denval q_prob_xt(scl::particles* pf);
  scl::denval p_prob_xt(scl::particles* pf);
  scl::denval p_prob_yt(scl::particles* pf); 
  INTEGER get_xdim() {return xdim;}
  INTEGER get_xlag() {return xlag;}
  INTEGER get_ydim() {return ydim;}
  INTEGER get_ylag() {return ylag;}
  sample draw_sample(INTEGER T, INT_32BIT& seed);
};


int main(int argc, char** argp, char** envp)
{
  #if defined DEBUG_CHECKS
    INTEGER T = 4;
    INTEGER N = 10;
  #else
    INTEGER T = 200;
    INTEGER N = 5000;
  #endif

  pfmod pfm;

  INT_32BIT seed = 780695;
  sample s = pfm.draw_sample(T, seed);

  INTEGER xdim = pfm.get_xdim();
  INTEGER xlag = pfm.get_xlag();
  INTEGER ydim = pfm.get_ydim();
  INTEGER ylag = pfm.get_ylag();

  realmat yplot(1,1+ylag+T);
  realmat xplot(2,1+xlag+T);
  xplot(1,1) = s.X0_Y0[1];
  xplot(2,1) = s.X0_Y0[2];
  yplot(1,1) = s.X0_Y0[3];
  xplot(1,2) = s.X0_Y0[4];
  xplot(2,2) = s.X0_Y0[5];
  yplot(1,2) = s.X0_Y0[6];
  for (INTEGER t=1; t<=T; ++t) {
    xplot(1,1+xlag+t) = s.X(1,t);
    xplot(2,1+xlag+t) = s.X(2,t);
    yplot(1,1+ylag+t) = s.Y(1,t);
  }
  

  #if defined DEBUG_CHECKS
    cout << '\n';
    cout << "xdim, xlag, ydim, ylag, T, N = " << xdim <<", "<< xlag 
         <<", "<< ydim <<", "<< ylag <<", "<< T <<", "<< N << '\n';
    cout << '\n';
    cout << "(x_{-xlag}, x_0, x_1, ..., x_T) = ";
    cout << xplot << '\n';
    cout << '\n';
    cout << "(y_{-ylag}, y_0, y_1, ..., y_T) = ";
    cout << yplot << '\n';
    cout << '\n';
  #endif

  particles pf(pfm,N,T);

  realmat ypf(1,T);
  for (INTEGER t=1; t<=T; ++t) ypf[t] = yplot[t+1]; 
  pf.set_Y(ypf);

  pf.generate_unconditional();

  #if defined DEBUG_CHECKS
    cout << "A on return from generate_unconditional" << '\n';
    for (INTEGER k=0; k<N; ++k) {
      intvec A_k = pf.get_A_k(k);
      for (INTEGER t=-pf.get_xlag(); t<=pf.get_T(); ++t) {
        cout << fmt('d',3,A_k[1+pf.get_xlag()+t]) <<' ';
      }
      cout << '\n';
    }
    cout << '\n';
    cout << "B on return from generate_unconditional" << '\n';
    for (INTEGER k=0; k<N; ++k) {
      intvec B_k = pf.get_B_k(k);
      for (INTEGER t=-pf.get_xlag(); t<=pf.get_T(); ++t) {
        cout << fmt('d',3,B_k[1+pf.get_xlag()+t]) <<' ';
      }
      cout << '\n';
    }
    cout << '\n';
    cout << '\n';
    cout << "XF on return from generate_unconditional" << '\n';
    cout << pf.get_XF();
    cout << '\n';
    cout << '\n';
    cout << "XS on return from generate_unconditional" << '\n';
    cout << pf.get_XS();
    cout << '\n';
    cout << '\n';
    cout << "YF on return from generate_unconditional" << '\n';
    cout << pf.get_YF();
    cout << '\n';
    cout << '\n';
    cout << "YS on return from generate_unconditional" << '\n';
    cout << pf.get_YS();
    cout << '\n';
    cout << '\n';
    cout << "W on return from generate_unconditional" << '\n';
    cout << pf.get_W();
    cout << '\n';
    cout << '\n';
    denval Z = pf.get_Z();
    cout << "Z = "<<boolalpha<<"("<< Z.positive <<", "<< Z.log_den <<")"<<'\n';
    cout << '\n';
  #endif

  realmat mean(xdim,1+xlag+T,0.0);
  for (INTEGER k=0; k<N; ++k) {
    for (INTEGER t=-xlag; t<=T; ++t) {
      for (INTEGER i=1; i<=xdim; ++i) {
        mean(i,1+xlag+t) += pf.x_s(k,i,t);
      }
    }
  }
  mean = mean/N;

  realmat sdev(xdim,1+xlag+T,0.0);
  for (INTEGER k=0; k<N; ++k) {
    for (INTEGER t=-xlag; t<=T; ++t) {
      for (INTEGER i=1; i<=xdim; ++i) {
        sdev(i,1+xlag+t) += pow(pf.x_s(k,i,t) - mean(i,1+xlag+t),2);
      }
    }
  }

  for (INTEGER t=-xlag; t<=T; ++t) {
    for (INTEGER i=1; i<=xdim; ++i) {
      sdev(i,1+xlag+t) = sqrt(sdev(i,1+xlag+t)/REAL(N-1));
    }
  }
  
   INTEGER kc =  pf.draw_k();

   #if defined DEBUG_CHECKS
     cout << "kc = " << kc << '\n';
     cout << '\n';
   #endif

   pf.generate_conditional(kc);

   #if defined DEBUG_CHECKS
    cout << "A on return from generate_conditional" << '\n';
    for (INTEGER k=0; k<N; ++k) {
      intvec A_k = pf.get_A_k(k);
      for (INTEGER t=-pf.get_xlag(); t<=pf.get_T(); ++t) {
        cout << fmt('d',3,A_k[1+pf.get_xlag()+t]) <<' ';
      }
      cout << '\n';
    }
    cout << '\n';
    cout << "B on return from generate_conditional" << '\n';
    for (INTEGER k=0; k<N; ++k) {
      intvec B_k = pf.get_B_k(k);
      for (INTEGER t=-pf.get_xlag(); t<=pf.get_T(); ++t) {
        cout << fmt('d',3,B_k[1+pf.get_xlag()+t]) <<' ';
      }
      cout << '\n';
    }
    cout << '\n';
    cout << '\n';
    cout << '\n';
    cout << "XF on return from generate_conditional" << '\n';
    cout << pf.get_XF();
    cout << '\n';
    cout << '\n';
    cout << "XS on return from generate_conditional" << '\n';
    cout << pf.get_XS();
    cout << '\n';
    cout << '\n';
    cout << "YF on return from generate_conditional" << '\n';
    cout << pf.get_YF();
    cout << '\n';
    cout << '\n';
    cout << "YS on return from generate_conditional" << '\n';
    cout << pf.get_YS();
    cout << '\n';
    cout << '\n';
    cout << "W on return from generate_conditional" << '\n';
    cout << pf.get_W();
    cout << '\n';
    cout << '\n';
    Z = pf.get_Z();
    cout << "Z = "<<boolalpha<<"("<< Z.positive <<", "<< Z.log_den <<")"<<'\n';
    cout << '\n';
  #endif

  ofstream fout;
  string filename = "tstpf.csv";
  fout.open(filename.c_str());
  if (!fout) error("Error, cannot open " + filename);

  for (INTEGER i=1; i<=xdim; ++i) fout << "mean" << i << ',';
  for (INTEGER i=1; i<=xdim; ++i) fout << "sdev" << i << ',';
  for (INTEGER i=1; i<=xdim; ++i) fout << "x" << i << ',';
  for (INTEGER i=1; i<=ydim-1; ++i) fout << "y" << i << ',';
  fout << "y" << ydim << '\n';

  for (INTEGER t=1; t<=1+xlag+T; ++t) {
    for (INTEGER i=1; i<=xdim; ++i) fout << mean(i,t) << ',';
    for (INTEGER i=1; i<=xdim; ++i) fout << sdev(i,t) << ',';
    for (INTEGER i=1; i<=xdim; ++i) fout << xplot(i,t)  << ',';
  for (INTEGER i=1; i<=ydim-1; ++i) fout << yplot(i,t) << ',';
  fout << yplot(ydim,t) << '\n';
  }

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

  return 0;
}


namespace {

  const INTEGER spin = 50;

  const REAL top =  230.0;
  const REAL bot = -230.0;

  const REAL etop = exp(top);
  const REAL ebot = exp(bot);

  REAL texp(REAL x)
  {
    if (x > top) {
      return etop;
    }
    else if (x < bot) {
      return ebot; 
    }
    else {
      return exp(x);
    }
  }

  const REAL logtwopi = log(6.283195307179587);

  const INT_32BIT fixed_seed = 5377065;
}

void pfmod::set_parms(const realmat& theta)
{
  if(theta.size() != 14) error("Error, pfmod, wrong dimension for theta");
  a0  = theta[1];
  a1  = theta[2];
  b01 = theta[3];
  b02 = theta[4];
  b11 = theta[5];
  b21 = theta[6];
  b12 = theta[7];
  b22 = theta[8];
  r11 = theta[9];
  r21 = theta[10];
  r22 = theta[11];
  r31 = theta[12];
  r32 = theta[13];
  r33 = theta[14];
  INT_32BIT seed = fixed_seed;
  sample s = draw_sample(1, seed);
  Mean_X0_Y0 = s.Mean_X0_Y0;
  Var_X0_Y0 = s.Var_X0_Y0;
}

realmat pfmod::get_parms()
{
  realmat theta(14,1);
  theta[1] = a0;
  theta[2] = a1;
  theta[3] = b01;
  theta[4] = b02;
  theta[5] = b11;
  theta[6] = b21;
  theta[7] = b12;
  theta[8] = b22;
  theta[9] = r11;
  theta[10] = r21;
  theta[11] = r22;
  theta[12] = r31;
  theta[13] = r32;
  theta[14] = r33;
  return theta;
}

void pfmod::q_draw_x0_y0(scl::particles* pf, INT_32BIT& seed)
{
  sample s = draw_sample(spin,seed);
  const INTEGER t = 0;
  pf->x(1,t) = s.X(1,spin);
  pf->x(2,t) = s.X(2,spin);
  pf->y(1,t) = s.Y(1,spin);
  pf->x(1,t-1) = s.X(1,spin-1);
  pf->x(2,t-1) = s.X(2,spin-1);
  pf->y(1,t-1) = s.Y(1,spin-1);
}

denval pfmod::q_prob_x0_y0(scl::particles* pf) {
  return psave;
}

denval pfmod::p_prob_x0_y0(scl::particles* pf)
{
  const INTEGER m = 2*(xdim + ydim);
  const INTEGER t = 0;
  realmat u(m,1);
  u[1] = pf->x(1,t) - Mean_X0_Y0[1];
  u[2] = pf->x(2,t) - Mean_X0_Y0[2];
  u[3] = pf->y(1,t) - Mean_X0_Y0[3];
  u[4] = pf->x(1,t-1) - Mean_X0_Y0[4];
  u[5] = pf->x(2,t-1) - Mean_X0_Y0[5];
  u[6] = pf->y(1,t-1) - Mean_X0_Y0[6];
  realmat R;
  cholesky(Var_X0_Y0,R);
  REAL logdetV = 2.0*logdetR(R);
  realmat Rinv;
  rinv(R,Rinv);
  realmat z = T(Rinv)*u;
  REAL log_prob = -0.5*REAL(m)*logtwopi - 0.5*logdetV - 0.5*(T(z)*z)[1];
  denval rv(true,log_prob);
  psave = rv;
  return rv;
}

void pfmod::q_draw_xt(scl::particles* pf, INT_32BIT& seed)
{
  const INTEGER t = 0;
  REAL e1 = unsk(seed);
  REAL e2 = unsk(seed);
  pf->x(1,t) = b01 + b11*pf->x(1,t-1) + b12*pf->x(2,t-1) + r11*e1;
  pf->x(2,t) = b02 + b21*pf->x(1,t-1) + b22*pf->x(2,t-1) + r21*e1 + r22*e2;
}

denval pfmod::q_prob_xt(scl::particles* pf)
{
  return psave;
}

denval pfmod::p_prob_xt(scl::particles* pf)
{
  const INTEGER m = 2;
  const INTEGER t = 0;
  realmat u(m,1);
  u[1] = pf->x(1,t) - b01 - b11*pf->x(1,t-1) - b12*pf->x(2,t-1);
  u[2] = pf->x(2,t) - b02 - b21*pf->x(1,t-1) - b22*pf->x(2,t-1);
  realmat R(m,m);
  R(1,1) = r11;
  R(1,2) = 0.0;
  R(2,1) = r21;
  R(2,2) = r22;
  REAL logdetV = 2.0*(log(r11)+log(r22));
  realmat Rinv = inv(R);
  realmat z = Rinv*u;
  REAL log_prob = -0.5*REAL(m)*logtwopi - 0.5*logdetV - 0.5*(T(z)*z)[1];
  denval rv(true,log_prob);
  psave = rv;
  return rv;
}

denval pfmod::p_prob_yt(scl::particles* pf)
{
  const INTEGER m = 2;
  const INTEGER t = 0;
  realmat e(2,1);
  e[1] = pf->x(1,t) - b01 - b11*pf->x(1,t-1) - b12*pf->x(2,t-1);
  e[2] = pf->x(2,t) - b02 - b21*pf->x(1,t-1) - b22*pf->x(2,t-1);
  realmat R(m,m);
  R(1,1) = r11;
  R(1,2) = 0.0;
  R(2,1) = r21;
  R(2,2) = r22;
  realmat Rinv = inv(R);
  e = Rinv*e;
  REAL mu = a0 + a1*pf->y(1,t-1) 
            + texp(pf->x(1,t-1) + pf->x(2,t-1))*(r31*e[1] + r32*e[2]);
  REAL sd = texp(pf->x(1,t-1)+pf->x(2,t-1))*r33;
  REAL z = (pf->y(1,t) - mu)/sd;
  REAL log_prob = -0.5*logtwopi - log(sd) - 0.5*z*z;
  denval rv(true,log_prob);
  psave = rv;
  return rv;
}

sample pfmod::draw_sample(INTEGER n, INT_32BIT& seed)
{
  const INTEGER m = 2*(xdim + ydim);
  sample s;
  s.X.resize(xdim,n);
  s.Y.resize(ydim,n);
  s.X0_Y0.resize(m,1);
  s.Mean_X0_Y0.resize(m,1,0.0);
  s.Var_X0_Y0.resize(m,m,0.0);
  realmat x(xdim,1);
  realmat xlag(xdim,1,0.0);
  REAL y;
  REAL ylag = 0.0;
  realmat xy(m,spin);
  INT_32BIT spin_seed = fixed_seed;
  for (INTEGER t=1; t<=spin; ++t) {
    REAL e1 = unsk(spin_seed);
    REAL e2 = unsk(spin_seed);
    REAL e3 = unsk(spin_seed);
    x[1] = b01 + b11*xlag[1] + b12*xlag[2] + r11*e1;
    x[2] = b02 + b21*xlag[1] + b22*xlag[2] + r21*e1 + r22*e2;
    REAL u = r31*e1 + r32*e2 + r33*e3;
    y = a0 + a1*ylag + texp(xlag[1]+xlag[2])*u;
    s.Mean_X0_Y0[4] += s.X0_Y0[4] = xlag[1];
    s.Mean_X0_Y0[5] += s.X0_Y0[5] = xlag[2];
    s.Mean_X0_Y0[6] += s.X0_Y0[6] = ylag;
    s.Mean_X0_Y0[1] += s.X0_Y0[1] = xy(1,t) = xlag[1] = x[1];
    s.Mean_X0_Y0[2] += s.X0_Y0[2] = xy(2,t) = xlag[2] = x[2];
    s.Mean_X0_Y0[3] += s.X0_Y0[3] = xy(3,t) = ylag = y;
  }
  s.Mean_X0_Y0 = s.Mean_X0_Y0/spin;
  for (INTEGER t=1; t<=spin; ++t) {
    for (INTEGER j=1; j<=m; ++j) {
      for (INTEGER i=1; i<=m; ++i) {
        s.Var_X0_Y0(i,j) 
	  += (xy(i,t) - s.Mean_X0_Y0[i]) * (xy(j,t) - s.Mean_X0_Y0[j]); 
      }
    }
  }
  s.Var_X0_Y0 = s.Var_X0_Y0/spin;
  for (INTEGER t=1; t<=n; ++t) {
    REAL e1 = unsk(seed);
    REAL e2 = unsk(seed);
    REAL e3 = unsk(seed);
    x[1] = b01 + b11*xlag[1] + b12*xlag[2] + r11*e1;
    x[2] = b02 + b21*xlag[1] + b22*xlag[2] + r21*e1 + r22*e2;
    REAL u = r31*e1 + r32*e2 + r33*e3;
    y = a0 + a1*ylag + texp(xlag[1]+xlag[2])*u;
    s.X(1,t) = xlag[1] = x[1];
    s.X(2,t) = xlag[2] = x[2];
    s.Y[t] = ylag = y;
  }
  return s;
}
