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

Copyright (C) 2004, 2005, 2006, 2007.

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.

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

#include <sstream>
#include <cerrno>
#include "libsnp.h"
#include "libsmm.h"
#include "snpusr.h"
#include "snp.h"
#include "emm.h"

using namespace scl;
using namespace libsnp;
using namespace libsmm;
using namespace snp;
using namespace emm;
using namespace std;

namespace {
  
  bool good_sim(const realmat& y)
  {
    #if defined GNU_GPP_COMPILER
      for (INTEGER i=1; i<=y.size(); ++i) {
        if (!finite(y[i])) return false;
      }
    #else
      const REAL quarter_max = REAL_MAX/4.0;
      for (INTEGER i=1; i<=y.size(); ++i) {
        if (fabs(y[i])>quarter_max) return false;
      }
    #endif
    return true;
  }

}

emm::snp_usrmod::snp_usrmod
  (const realmat& data, INTEGER len_mod_parm, INTEGER len_mod_func,
   const std::vector<std::string>& mod_pfvec,
   const std::vector<std::string>& mod_alvec, 
   std::ostream& detail)
{
  parmfile pf;

  if(!pf.set_parms(mod_pfvec,detail)) {
     detail.flush();
     string msg("Error, snp_usrmod, cannot read parmfile.  Make sure\n");
     msg += "the typedef in emmusr.h is bound to the correct class.";
     error(msg);
  }
    
  datread_type dr; 
  dr.initialize(pf.get_datparms());
  
  Y = data;
  if((Y.get_rows()!=pf.get_datparms().M)||(Y.get_cols()!=pf.get_datparms().n)){ 
    error("Error, snp_usrmod, data rows or cols differ from snp parmfile");
  }

  if (pf.get_optparms().print) {
    detail << starbox("/First 12 observations//");
    detail << Y("",seq(1,12));
    detail << starbox("/Last 12 observations//");
    detail << Y("",seq(pf.get_datparms().n-11, pf.get_datparms().n));
  }

  trnfrm tr_tmp(pf.get_datparms(), pf.get_tranparms(), Y);
  tr = tr_tmp;

  pf.set_tranparms(tr.get_tranparms());
    
  if (pf.get_optparms().print) {
    detail << starbox("/Mean and variance of data//");
    if (pf.get_tranparms().diag) {
      detail << "(Variance has been diagonalized.)\n";
    }
    detail << pf.get_tranparms().mean << pf.get_tranparms().variance;
  }

  X = Y;

  tr.normalize(Y);
  tr.normalize(X);

  if (pf.get_tranparms().squash == 1) {
    tr.spline(X);
  } else if (pf.get_tranparms().squash == 2) {
    tr.logistic(X);
  }

  if (pf.get_optparms().print) {
    detail << starbox("/First 12 normalized observations//");
    detail << Y("",seq(1,12));
    detail << starbox("/Last 12 normalized observations//");
    detail << Y("",seq(pf.get_datparms().n-11, pf.get_datparms().n));
    if (pf.get_tranparms().squash > 0) {
      detail << starbox("/First 12 transformed observations//");
      detail << Y("",seq(1,12));
      detail << starbox("/Last 12 transformed observations//");
      detail << Y("",seq(pf.get_datparms().n-11, pf.get_datparms().n));
    }
  }

  if (pf.get_optparms().print) detail.flush();

  snpll ll_tmp(pf.get_datparms(),
    pf.get_snpden(), pf.get_afunc(), pf.get_ufunc(),
    pf.get_rfunc(), pf.get_afunc_mask(), pf.get_ufunc_mask(),
    pf.get_rfunc_mask());

  ll = ll_tmp;
  ll.set_XY(&X,&Y);

  vector<string>::const_iterator u_ptr=mod_alvec.begin();
  
  seed = pf.get_optparms().iseed;
  ++u_ptr;
  slen = atoi((u_ptr++)->substr(0,12).c_str());
  spin = atoi((u_ptr++)->substr(0,12).c_str());
  lrho = ll.get_rho().size();
  if (lrho != len_mod_parm) {
    error("Error, snp_usrmod, number of parameters differs from snp parmfile");
  }
  lstats = pf.get_datparms().M;
  if (lstats != len_mod_func) {
    error("Error, snp_usrmod, number of functionals differs from snp_usrmod");
  }

  if (pf.get_optparms().print) {
    detail << '\n';
    detail << "\t seed = " << seed << '\n';
    detail << "\t spin = " << spin << '\n';
    detail << "\t slen = " << slen << '\n';
    detail << "\t lrho = " << lrho << '\n';
    detail << "\t lstats = " << lstats << '\n';
    detail.flush();
  }
}

emm::snp_usrmod::snp_usrmod() 
: tr(), ll(), X(), Y(), seed(), spin(), slen(), lrho(), lstats() 
{ }

emm::snp_usrmod::snp_usrmod(const snp_usrmod& obj)
: tr(obj.tr), ll(obj.ll), X(obj.X), Y(obj.Y), 
  seed(obj.seed), spin(obj.spin), slen(obj.slen), 
  lrho(obj.lrho), lstats(obj.lstats) 
{ }

emm::snp_usrmod& snp_usrmod::operator=(const snp_usrmod& obj)
{
  if (this != &obj) {
    tr = obj.tr; ll = obj.ll; X = obj.X; Y = obj.Y; 
    seed = obj.seed; spin = obj.spin; slen = obj.slen; 
    lrho = obj.lrho; lstats = obj.lstats;
  }
  return *this;
}

bool emm::snp_usrmod::snp_simulate(INT_32BIT& simseed, realmat& sim)
{
  if (&Y==0||&X==0) error("Error, simulate, data not initialized");

  datparms dpm = ll.get_datparms(); 
  tranparms tpm = tr.get_tranparms();

  snpden f = ll.get_snpden();
  afunc af = ll.get_afunc();
  ufunc uf = ll.get_ufunc();
  rfunc rf = ll.get_rfunc();

  if (dpm.M != Y.nrow() || dpm.M != X.nrow() || dpm.M != f.get_ly()) {
    error("Error, snp_simulate, this should never happen");
  }

  if (sim.get_rows() != dpm.M || sim.get_cols() != slen) {
    sim.resize(dpm.M,slen);
  }

  realmat dawa0,dawA;
  realmat duwb0, duwb0_lag;
  kronprd duwB, duwB_lag;
  realmat dRwb0,dRwB;
  realmat dRwRparms;

  realmat y(dpm.M,1);
  realmat x(dpm.M,1);
  realmat u(dpm.M,1);
  realmat y_lag(dpm.M,1);
  realmat x_lag(dpm.M,1);
  realmat u_lag(dpm.M,1);

  af.initialize_state();
  uf.initialize_state();
  rf.initialize_state();

  for (INTEGER i=1; i<=dpm.M; ++i) {
    x[i] = X(i,1);
  }

  u = uf(x,duwb0,duwB);

  for (INTEGER t=3; t<=dpm.drop; ++t) {

    for (INTEGER i=1; i<=dpm.M; ++i) {
      y[i] = Y(i,t);
      x[i] = X(i,t);
      y_lag[i] = Y(i,t-1);
      x_lag[i] = X(i,t-1);
    }

    u_lag = u;
    duwb0_lag = duwb0;
    duwB_lag = duwB;

    u = uf(x_lag,duwb0,duwB);

    f.set_R(rf(y_lag,u_lag,x_lag,duwb0_lag,duwB_lag,dRwb0,dRwB,dRwRparms));
    f.set_a(af(x_lag,dawa0,dawA));
    f.set_u(u);
  }

  for (INTEGER t=dpm.drop+1; t<=spin; ++t) {

    for (INTEGER i=1; i<=dpm.M; ++i) {
      y_lag[i] = y[i];
      x_lag[i] = x[i];
    }

    u_lag = u;
    duwb0_lag = duwb0;
    duwB_lag = duwB;

    u = uf(x_lag,duwb0,duwB);

    f.set_R(rf(y_lag,u_lag,x_lag,duwb0_lag,duwB_lag,dRwb0,dRwB,dRwRparms));
    f.set_a(af(x_lag,dawa0,dawA));
    f.set_u(u);

    y = f.sampy(simseed);
    if (!good_sim(y)) return false;
    x = y;
    if (tpm.squash == 1) {
      tr.spline(x);
    }
    else if (tpm.squash == 2) {
      tr.logistic(x);
    }
  }

  for (INTEGER t=1; t<=slen; ++t) {

    for (INTEGER i=1; i<=dpm.M; ++i) {
      y_lag[i] = y[i];
      x_lag[i] = x[i];
    }

    u_lag = u;
    duwb0_lag = duwb0;
    duwB_lag = duwB;

    u = uf(x_lag,duwb0,duwB);

    f.set_R(rf(y_lag,u_lag,x_lag,duwb0_lag,duwB_lag,dRwb0,dRwB,dRwRparms));
    f.set_a(af(x_lag,dawa0,dawA));
    f.set_u(u);

    y = f.sampy(simseed);
    if (!good_sim(y)) return false;
    x = y;
    if (tpm.squash == 1) {
      tr.spline(x);
    }
    else if (tpm.squash == 2) {
      tr.logistic(x);
    }

    realmat v = y;
    tr.unnormalize(v);

    for (INTEGER i=1; i<=v.size(); ++i) {
      sim(i,t) = v[i];
    }
  }

  return true;
}

bool emm::snp_usrmod::gen_sim(realmat& sim) 
{
  INT_32BIT fixed_seed = 770116;
  return snp_simulate(fixed_seed,sim);
} 

bool emm::snp_usrmod::gen_sim(realmat& sim, realmat& stats) 
{ 
  INT_32BIT fixed_seed = 770116;
  if (!snp_simulate(fixed_seed,sim)) return false;
  realmat sum(sim.get_rows(),1,0.0);
  for (INTEGER t=1; t<=sim.get_cols(); ++t) sum += sim(intvec(),t);
  stats = sum/sim.get_cols();
  return true;
}

bool emm::snp_usrmod::gen_bootstrap(vector<realmat>& bs) 
{
  realmat sim;

  if (!snp_simulate(seed,sim)) return false; 

  INTEGER n_sim = sim.get_cols();
  INTEGER n_dat = Y.get_cols();

  if (n_sim < n_dat) return false;

  INTEGER len_vec = 2*lrho+1;
  INTEGER inc, len;

  if (n_sim/n_dat < len_vec) {
    len = n_sim/n_dat;
    inc = n_dat;
  }
  else {
    len = len_vec;
    inc = n_sim/len;
  }

  len_vec = bs.size();

  if (len_vec != len) {
    bs.resize(len);
  }

  INTEGER s = 0;
  for (INTEGER i=0; i<len; i++) {
    bs[i] = sim("",seq(s+1,s+n_dat));
    s+=inc;
  }

  return true;
}

void emm::snp_usrmod::get_rho(realmat& rho) 
{ 
  rho = ll.get_rho();
}

void emm::snp_usrmod::set_rho(const realmat& rho) 
{ 
  ll.set_rho(rho);
}

bool emm::snp_usrmod::support(const realmat& rho) 
{
  realmat test = ll.get_theta()(ll.get_srvec(),1);
  for (INTEGER i=1; i<=test.size(); ++i) if (test[i] <= 0.0) return false;
  return true; 
}

den_val emm::snp_usrmod::prior(const realmat& rho, const realmat& stats) 
{
  return den_val(true,0.0); 
}
 
