/*-----------------------------------------------------------------------------

Copyright (C) 2014

A. Ronald Gallant
Post Office Box 659
Chapel Hill NC 27514-0659
USA   

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

-------------------------------------------------------------------------------

This is part of an implementation of the generic sequential Monte Carlo 
algorithm and the conditional sequential Monte Carlo update described 
in Andrieu, C., A. Douced, and R. Holenstein (2010), "Particle Markov Chain 
Monte Carlo Methods," Journal of the Royal Statistical Society, Series B, 
72, 269--342.

Usage is described in the header particles.h.

-----------------------------------------------------------------------------*/

#include "libscl.h"
#include "pfmod.h"

using namespace std;
using namespace scl;

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;
}
