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

Copyright (C) 2004, 2006, 2009, 2014, 2019.

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 "snp.h"

using namespace scl;
using namespace libsnp;
using namespace snp;
using namespace std;

snp::snpll::snpll()
: dpm(datparms()), Y(0), X(0), f(snpden()), 
  af(afunc()), uf(ufunc()), rf(rfunc()), 
  af_mask(afunc()), uf_mask(ufunc()), rf_mask(rfunc()),
  rho(realmat()), theta(realmat()), srvec(intvec()), 
  rt(vector< pair<INTEGER,INTEGER> >()),
  a0(realmat()), A(realmat()), b0(realmat()), B(realmat()), 
  Rparms(realmat()),
  compute_infmat(false), infmat(realmat()),
  compute_scores(false), scores(realmat())
{ }

snp::snpll::snpll(datparms dp,
      libsnp::snpden fn, libsnp::afunc afn,
      libsnp::ufunc ufn, libsnp::rfunc rfn,
      libsnp::afunc afn_mask,
      libsnp::ufunc ufn_mask, libsnp::rfunc rfn_mask)
: dpm(dp), Y(0), X(0), f(fn), 
  af(afn), uf(ufn), rf(rfn), 
  af_mask(afn_mask), uf_mask(ufn_mask), rf_mask(rfn_mask),
  rho(realmat()), theta(realmat()), srvec(intvec()), 
  rt(vector< pair<INTEGER,INTEGER> >()),
  a0(afn.get_a0()), A(afn.get_A()), b0(ufn.get_b0()), B(ufn.get_B()), 
  Rparms(rfn.get_Rparms()),
  compute_infmat(false), infmat(realmat()),
  compute_scores(false), scores(realmat())
{ 

  INTEGER ltheta = a0.size()+A.size()+b0.size()+B.size()+Rparms.size();
  theta.resize(ltheta,1);

  realmat a0_mask = af_mask.get_a0(); //In these, -1 is fixed, and 0 is active.
  realmat A_mask = af_mask.get_A();   //Add 1 to get the standard convention.
  A_mask[1] = -1.0; 
  realmat b0_mask = uf_mask.get_b0();
  realmat B_mask = uf_mask.get_B();
  realmat Rparms_mask = rf_mask.get_Rparms();

  intvec mask(ltheta,0);

  INTEGER count=0;
  INTEGER lrho = 0;

  for (INTEGER i=1; i<=a0.size(); ++i) {
     ++count;
     theta[count] = a0[i];
     mask[count] = INTEGER(a0_mask[i] + 1.0);  //Note mask coding convention
     lrho += mask[count];                      //of -1, 0 described above.
  }

  for (INTEGER i=1; i<=A.size(); ++i) {
     ++count;
     theta[count] = A[i];
     mask[count] = INTEGER(A_mask[i] + 1.0); 
     lrho += mask[count];
  }

  for (INTEGER i=1; i<=b0.size(); ++i) {
     ++count;
     theta[count] = b0[i];
     mask[count] = INTEGER(b0_mask[i] + 1.0); 
     lrho += mask[count];
  }

  for (INTEGER i=1; i<=B.size(); ++i) {
     ++count;
     theta[count] = B[i];
     mask[count] = INTEGER(B_mask[i] + 1.0); 
     lrho += mask[count];
  }

  for (INTEGER i=1; i<=Rparms.size(); ++i) {
     ++count;
     theta[count] = Rparms[i];
     mask[count] = INTEGER(Rparms_mask[i] + 1.0); 
     lrho += mask[count];
  }

  if (count != ltheta) error("Error, snpll, this is inexplicable");

  rho.resize(lrho,1);
  rt.resize(lrho);
 
  count = 0;
  for (INTEGER i=1; i<=ltheta; ++i) {
    if (mask[i]) {
      count++;
      rt[count-1].first = count;
      rt[count-1].second = i;
    }
  }

  vector< pair<INTEGER,INTEGER> >::const_iterator rtptr; 

  for (rtptr = rt.begin(); rtptr != rt.end(); ++rtptr) {
     rho[rtptr->first] = theta[rtptr->second];
  }

  intvec srvec_elements(ltheta);

  INTEGER theta_count=0;

  for (INTEGER i=1; i<=af.get_nrowa0(); ++i) {
    ++theta_count;
  }
  for (INTEGER j=1; j<=af.get_ncolA(); ++j) {
    for (INTEGER i=1; i<=af.get_nrowA(); ++i) {
      ++theta_count;
    }
  }
  if (uf.is_intercept()) {
    for (INTEGER i=1; i<=uf.get_ly(); ++i) {
      ++theta_count;
    }
  }
  if (uf.is_regression()) {
    for (INTEGER j=1; j<=uf.get_lx()*uf.get_lag(); ++j) {
      for (INTEGER i=1; i<=uf.get_ly(); ++i) {
        ++theta_count;
      }
    }
  }

  INTEGER srvec_count = 0;
  
  for (INTEGER j=1; j<=rf.get_ly(); ++j) {
    for (INTEGER i=1; i<=j; ++i) {
      ++theta_count;
      if (i == j) {
        srvec_elements[++srvec_count] = theta_count;
      }
    }
  }
  for (INTEGER j=1; j<=rf.get_colsP(); ++j) {
    for (INTEGER i=1; i<=rf.get_rowsP(); ++i) {
      ++theta_count;
      char t = rf.get_Ptype();
      INTEGER k = j%rf.get_rowsP();
      if ((t=='s') || (t=='d' && i==1) || (t=='f' && i==1 && k==1)) {
        srvec_elements[++srvec_count] = theta_count;
      }
    }
  }
  for (INTEGER j=1; j<=rf.get_colsQ(); ++j) {
    for (INTEGER i=1; i<=rf.get_rowsQ(); ++i) {
      ++theta_count;
      char t = rf.get_Qtype();
      INTEGER k = j%rf.get_rowsQ();
      if ((t=='s') || (t=='d' && i==1) || (t=='f' && i==1 && k==1)) {
        srvec_elements[++srvec_count] = theta_count;
      }
    }
  }
  for (INTEGER j=1; j<=rf.get_colsV(); ++j) {
    for (INTEGER i=1; i<=rf.get_rowsV(); ++i) {
      ++theta_count;
      /*
      char t = rf.get_Vtype();
      INTEGER k = j%rf.get_rowsV();
      if ((t=='s') || (t=='d' && i==1) || (t=='f' && i==1 && k==1)) {
        srvec_elements[++srvec_count] = theta_count;
      }
      */
    }
  }
  for (INTEGER j=1; j<=rf.get_colsW(); ++j) {
    for (INTEGER i=1; i<=rf.get_rowsW(); ++i) {
      ++theta_count;
      char t = rf.get_Wtype();
      INTEGER k = j%rf.get_rowsW();
      if ((t=='s') || (t=='d' && i==1) || (t=='f' && i==1 && k==1)) {
        srvec_elements[++srvec_count] = theta_count;
      }
    }
  }

  if (theta_count != ltheta) error("Error, snpll, this is inexplicable");

  srvec.resize(srvec_count);
  
  for (INTEGER i=1; i<=srvec_count; ++i) {
    srvec[i] = srvec_elements[i];
  }
}

snp::snpll::snpll(const snpll& sl)
: dpm(sl.dpm), Y(sl.Y), X(sl.X), 
  f(sl.f), af(sl.af), uf(sl.uf), rf(sl.rf), 
  af_mask(sl.af_mask), uf_mask(sl.uf_mask), rf_mask(sl.rf_mask),
  rho(sl.rho), theta(sl.theta), srvec(sl.srvec), rt(sl.rt),
  a0(sl.a0), A(sl.A), b0(sl.b0), B(sl.B), Rparms(sl.Rparms),
  compute_infmat(sl.compute_infmat), infmat(sl.infmat),
  compute_scores(sl.compute_scores), scores(sl.scores)
{ }

snpll& snp::snpll::operator=(const snpll& sl)
{
  if (this != &sl) {
    dpm=sl.dpm; Y=sl.Y; X=sl.X; 
    f=sl.f; af=sl.af; uf=sl.uf; rf=sl.rf; 
    af_mask=sl.af_mask; uf_mask=sl.uf_mask; rf_mask=sl.rf_mask;
    rho=sl.rho; theta=sl.theta; srvec = sl.srvec, rt=sl.rt;
    a0=sl.a0; A=sl.A; b0=sl.b0; B=sl.B; Rparms=sl.Rparms;
    compute_infmat = sl.compute_infmat; infmat = sl.infmat;
    compute_scores = sl.compute_scores; scores = sl.scores;
  }
  return *this;
}
  

void snp::snpll::agglomerate()
{
  INTEGER count = 0;

  for (INTEGER i=1; i<=a0.size(); ++i) {
    ++count;
    theta[count] = a0[i];
  }

  for (INTEGER i=1; i<=A.size(); ++i) {
    ++count;
    theta[count] = A[i];
  }

  for (INTEGER i=1; i<=b0.size(); ++i) {
    ++count;
    theta[count] = b0[i];
  }

  for (INTEGER i=1; i<=B.size(); ++i) {
    ++count;
    theta[count] = B[i];
  }


  for (INTEGER i=1; i<=Rparms.size(); ++i) {
    ++count;
    theta[count] = Rparms[i];
  }

  if (count != theta.size()) error("Error, snpll, this is inexplicable");

}

void snp::snpll::distribute()
{

  INTEGER count = 0;

  for (INTEGER i=1; i<=a0.size(); ++i) {
    ++count;
    a0[i] = theta[count];
  }

  for (INTEGER i=1; i<=A.size(); ++i) {
    ++count;
    A[i] = theta[count];
  }

  for (INTEGER i=1; i<=b0.size(); ++i) {
    ++count;
    b0[i] = theta[count];
  }


  for (INTEGER i=1; i<=B.size(); ++i) {
    ++count;
    B[i] = theta[count];
  }


  for (INTEGER i=1; i<=Rparms.size(); ++i) {
    ++count;
    Rparms[i] = theta[count];
  }

  if (count != theta.size()) error("Error, snpll, this is inexplicable");
}

void snp::snpll::set_XY(const realmat* x, const realmat* y)
{
  if (rho.size()<=0) error("Error, snpll, can't use if default construced");
  if (x->nrow() != y->nrow()) error("Error, snpll, row dim of x & y differ");
  if (x->ncol() != y->ncol()) error("Error, snpll, col dim of x & y differ");
  X = x;  Y = y;
}  

void snp::snpll::set_rho(const realmat& irho)
{
  INT_32BIT jseed = 740726;
  this->set_rho(irho, 0.0, 0.0, jseed);
}

void snp::snpll::set_rho(const realmat& irho, REAL fnew, REAL fold, 
        INT_32BIT& jseed)
{

  if (rho.size()<=0) error("Error, snpll, can't use if default construced");
  if (irho.size() != rho.size()) error("Error, snpll, length of irho wrong");

  rho = irho;

  vector< pair<INTEGER,INTEGER> >::const_iterator rtptr; 

  for (rtptr = rt.begin(); rtptr != rt.end(); ++rtptr) {
     theta[rtptr->second] = rho[rtptr->first];
  }

  distribute();

  if (fnew != 0.0 || fold != 0.0) {            //Note mask coding convention
    for (INTEGER i=1; i<=a0.size(); ++i) {     //of -1, 0 described above.
      if (af_mask.get_a0()[i] == 0.0) {
        REAL r = 2.0*ran(jseed) - 1.0;
        a0[i] += a0[i]==0.0 ? r*fnew : a0[i]*r*fold;
      }
    }
    for (INTEGER i=2; i<=A.size(); ++i) {
      if (af_mask.get_A()[i] == 0.0) {
        REAL r = 2.0*ran(jseed) - 1.0;
        A[i] += A[i]==0.0 ? r*fnew : A[i]*r*fold;
      }
    }
    fnew = fnew>0 ? fnew : 0.0;
    fold = fold>0 ? fold : 0.0;
    for (INTEGER i=1; i<=b0.size(); ++i) {
      if (uf_mask.get_b0()[i] == 0.0) {
        REAL r = 2.0*ran(jseed) - 1.0;
        b0[i] += b0[i]==0.0 ? r*fnew : b0[i]*r*fold;
      }
    }
    for (INTEGER i=1; i<=B.size(); ++i) {
      if (uf_mask.get_B()[i] == 0.0) {
        REAL r = 2.0*ran(jseed) - 1.0;
        B[i] += B[i]==0.0 ? r*fnew : B[i]*r*fold;
      }
    }
    for (INTEGER i=1; i<=Rparms.size(); ++i) {
      if (rf_mask.get_Rparms()[i] == 0.0) {
        REAL r = 2.0*ran(jseed) - 1.0;
        Rparms[i] += Rparms[i]==0.0 ? r*fnew : Rparms[i]*r*fold;
      }
    }
  }

  if (a0.size()>0) af.set_a0(a0);
  af.set_A(A);
  if (b0.size()>0) uf.set_b0(b0);
  if (B.size()>0) uf.set_B(B);
  rf.set_Rparms(Rparms);

  if (fnew != 0.0 || fold != 0.0) {
    agglomerate();
    for (rtptr = rt.begin(); rtptr != rt.end(); ++rtptr) {
       rho[rtptr->first] = theta[rtptr->second];
    }   
  }
}

REAL snp::snpll::log_likehood(realmat& dllwrho) 
{
  if (rho.size()<=0) error("Error, snpll, can't use if default construced");
  if (Y==0||X==0) error("Error, snpll, log_likehood, data not initialized");
  if (dpm.M != Y->nrow() || dpm.M != X->nrow() || dpm.M != f.get_ly())
    error("Error, snpll, this should never happen");

  // Errors found here should be corrected in 
  // REAL snpll::log_likehood(INTEGER& drop, realmat& llv, realmat& dllvwrho)
  // below as well

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

  realmat dlfwa, dlfwu, dlfwR;
  
  realmat dllwa0;
  if (a0.size()>0) dllwa0.resize(1,a0.size(),0.0);
  realmat dllwA(1,A.size(),0.0);
  realmat dllwb0;
  if (b0.size()>0) dllwb0.resize(1,b0.size(),0.0);
  realmat dllwB;
  if (B.size()>0) dllwB.resize(1,B.size(),0.0);
  realmat dllwRparms(1,Rparms.size(),0.0);

  REAL log_likelihood = 0.0;

  realmat theta_infmat;
  if (compute_infmat) theta_infmat.resize(theta.size(),theta.size(),0.0); 

  if (compute_scores) {
    if (!compute_infmat) {
      string msg = 
        "Error, snpll, compute_infmat must be true if compute_scores is true";
      error(msg);
    }
    scores.resize(rho.size(),dpm.n);
  }

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

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

    log_likelihood += f.log_f(y,dlfwa,dlfwu,dlfwR);

    if (a0.size()>0) dllwa0 += dlfwa*dawa0;
    dllwA += dlfwa*dawA;
    if (b0.size()>0) {
      dllwb0 += dlfwu*duwb0;
      if (dRwb0.size()>0) {
        dllwb0 += dlfwR*dRwb0;
      }
    }
    if (B.size()>0) {
      dllwB += dlfwu*duwB;
      if (dRwB.size()>0) {
        dllwB += dlfwR*dRwB;
      }
    }
    dllwRparms += dlfwR*dRwRparms;

    if (compute_infmat) {
      realmat dlfwtheta;
      if (a0.size()>0) {
        dlfwtheta = cbind(dlfwa*dawa0,dlfwa*dawA);
      }
      else {
        dlfwtheta = dlfwa*dawA;
      }
      if (b0.size()>0) {
        if (dRwb0.size()>0) {
          dlfwtheta = cbind(dlfwtheta,dlfwu*duwb0+dlfwR*dRwb0);
        }
        else {
          dlfwtheta = cbind(dlfwtheta,dlfwu*duwb0);
        }
      }
      if (B.size()>0) {
        if (dRwB.size()>0) {
          dlfwtheta = cbind(dlfwtheta,dlfwu*duwB+dlfwR*dRwB);
        }
        else {
          dlfwtheta = cbind(dlfwtheta,dlfwu*duwB);
        }
      }
      dlfwtheta = cbind(dlfwtheta,dlfwR*dRwRparms);

      for (INTEGER j=1; j<=dlfwtheta.ncol(); ++j) {
        for (INTEGER i=1; i<=dlfwtheta.ncol(); ++i) {
          theta_infmat(i,j) += dlfwtheta[i]*dlfwtheta[j];
        }
      }

      if (compute_scores) {
        vector< pair<INTEGER,INTEGER> >::const_iterator rti_ptr;
        if (t <= dpm.drop) {
          for (INTEGER i=1; i<=scores.nrow(); ++i) {
            scores(i,t) = 0.0;
          }
        }
        else {
          for (rti_ptr = rt.begin(); rti_ptr != rt.end(); ++rti_ptr) {
            scores(rti_ptr->first,t) = dlfwtheta[rti_ptr->second];
          }
        }
      }

    }
  }

  realmat dllwtheta; 
  if (a0.size()>0) {
    dllwtheta = cbind(dllwa0,dllwA);
  }
  else {
    dllwtheta = dllwA;
  }
  if (b0.size()>0) dllwtheta = cbind(dllwtheta,dllwb0);
  if (B.size()>0) dllwtheta = cbind(dllwtheta,dllwB);
  dllwtheta = cbind(dllwtheta,dllwRparms);

  dllwrho.resize(1,rho.size());

  vector< pair<INTEGER,INTEGER> >::const_iterator rti_ptr; 

  for (rti_ptr = rt.begin(); rti_ptr != rt.end(); ++rti_ptr) {
    dllwrho[rti_ptr->first] = dllwtheta[rti_ptr->second];
  }
  
  if (compute_infmat) {
    infmat.resize(rho.size(),rho.size()); 
    vector< pair<INTEGER,INTEGER> >::const_iterator rtj_ptr; 
    for (rti_ptr = rt.begin(); rti_ptr != rt.end(); ++rti_ptr) {
      for (rtj_ptr = rt.begin(); rtj_ptr != rt.end(); ++rtj_ptr) {
        infmat(rti_ptr->first,rtj_ptr->first) =
           theta_infmat(rti_ptr->second,rtj_ptr->second);
      }
    }
  }

  return log_likelihood;
}

REAL snp::snpll::log_likehood(realmat& dllwrho, realmat& rho_infmat) 
{
  compute_infmat = true;
  REAL ll = this->log_likehood(dllwrho);
  rho_infmat = infmat;
  compute_infmat = false;
  return ll;
}

REAL snp::snpll::log_likehood(realmat& dllwrho, realmat& rho_infmat,
  realmat& rho_scores) 
{
  compute_infmat = true;
  compute_scores = true;
  REAL ll = this->log_likehood(dllwrho);
  rho_infmat = infmat;
  rho_scores = scores;
  compute_infmat = false;
  compute_scores = false;
  return ll;
}

  
REAL snp::snpll::log_likehood(INTEGER& drop, realmat& llv, realmat& dllvwrho) 
{
  if (rho.size()<=0) error("Error, snpll, can't use if default construced");
  if (Y==0||X==0) error("Error, snpll, log_likehood, data not initialized");
  if (dpm.M != Y->nrow() || dpm.M != X->nrow() || dpm.M != f.get_ly())
    error("Error, snpll, this should never happen");

  // A modification of 
  // REAL snpll::log_likehood(realmat& dllwrho) 
  // Errors found here should be corrected there.

  llv.resize(dpm.n,1);
  dllvwrho.resize(dpm.n,rho.size());

  realmat dllvwtheta(dpm.n,theta.size());

  drop = dpm.drop;
  for (INTEGER t=1; t<=drop; ++t) {
    llv[t] = 0.0;
    for (INTEGER j=1; j<=theta.size(); ++j) {
      dllvwtheta(t,j) = 0.0;
    }
  }

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

  realmat dlfwa, dlfwu, dlfwR;
  
  realmat dllwa0;
  if (a0.size()>0) dllwa0.resize(1,a0.size(),0.0);
  realmat dllwA(1,A.size(),0.0);
  realmat dllwb0;
  if (b0.size()>0) dllwb0.resize(1,b0.size(),0.0);
  realmat dllwB;
  if (B.size()>0) dllwB.resize(1,B.size(),0.0);
  realmat dllwRparms(1,Rparms.size(),0.0);

  REAL log_likelihood = 0.0;

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

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

    REAL ll_t = f.log_f(y,dlfwa,dlfwu,dlfwR);
    llv[t] = ll_t;
    log_likelihood += ll_t;

    realmat tmp;
    INTEGER loc = 0;

    if (a0.size()>0) {
      tmp = dlfwa*dawa0;
      //dllwa0 += tmp;
      for (INTEGER j=1; j<=tmp.size(); ++j) {
        ++loc;
        dllvwtheta(t,loc) = tmp[j];
      }
    }

    tmp = dlfwa*dawA;
    //dllwA += tmp;
    for (INTEGER j=1; j<=tmp.size(); ++j) {
      ++loc;
      dllvwtheta(t,loc) = tmp[j];
    }
    
    if (b0.size()>0) {
      tmp = dlfwu*duwb0;
      //dllwb0 += tmp;
      for (INTEGER j=1; j<=tmp.size(); ++j) {
        ++loc;
        dllvwtheta(t,loc) = tmp[j];
      }
      if (dRwb0.size()>0) {
        tmp = dlfwR*dRwb0;
        //dllwb0 += tmp;
        for (INTEGER j=1; j<=tmp.size(); ++j) {
          ++loc;
          dllvwtheta(t,loc) = tmp[j];
        }
      }
    }

    if (B.size()>0) {
      tmp = dlfwu*duwB;
      //dllwB += tmp;
      for (INTEGER j=1; j<=tmp.size(); ++j) {
        ++loc;
        dllvwtheta(t,loc) = tmp[j];
      }
      if (dRwB.size()>0) {
        tmp = dlfwR*dRwB;
        //dllwB += tmp;
        for (INTEGER j=1; j<=tmp.size(); ++j) {
          ++loc;
          dllvwtheta(t,loc) = tmp[j];
        }
      }
    }

    tmp = dlfwR*dRwRparms;
    //dllwRparms += tmp;
    for (INTEGER j=1; j<=tmp.size(); ++j) {
      ++loc;
      dllvwtheta(t,loc) = tmp[j];
    }

    if (loc != theta.size()) error("Error, snpll, this should never happen");
  }

  vector< pair<INTEGER,INTEGER> >::const_iterator rti_ptr; 

  /*
  realmat dllwtheta; 
  if (a0.size()>0) {
    dllwtheta = cbind(dllwa0,dllwA);
  }
  else {
    dllwtheta = dllwA;
  }
  if (b0.size()>0) dllwtheta = cbind(dllwtheta,dllwb0);
  if (B.size()>0) dllwtheta = cbind(dllwtheta,dllwB);
  dllwtheta = cbind(dllwtheta,dllwRparms);

  dllwrho.resize(1,rho.size());

  for (rti_ptr = rt.begin(); rti_ptr != rt.end(); ++rti_ptr) {
    dllwrho[rti_ptr->first] = dllwtheta[rti_ptr->second];
  }
  */

  for (INTEGER t=1; t<=dpm.n; ++t) {
    for (rti_ptr = rt.begin(); rti_ptr != rt.end(); ++rti_ptr) {
      dllvwrho(t,rti_ptr->first) = dllvwtheta(t,rti_ptr->second);
    }
  }

  return log_likelihood;
}

ostream& snp::snpll::write_stats(std::ostream& os)
{
  os << "Fit criteria:\n";
  compute_infmat = true;
  realmat dllwrho;
  REAL ll = this->log_likehood(dllwrho);
  compute_infmat = false;
  os << "  Length rho =   " << fmt('d',10,rho.size()) << '\n';
  os << "  Length theta = " << fmt('d',10,theta.size()) << '\n';
  os << "  n - drop =     " << fmt('d',10,dpm.n-dpm.drop) << '\n';
  os << "  -2 ln likelihood = " << fmt('f',20,8,-2.0*ll)
           << fmt('e',27,17,-2.0*ll) << '\n';
  REAL sn = -ll/REAL(dpm.n);
  REAL aic = sn + (REAL(rho.size())/REAL(dpm.n));
  REAL hq =  sn + (REAL(rho.size())/REAL(dpm.n))*log(log(REAL(dpm.n)));
  REAL bic = sn + (REAL(rho.size())/REAL(dpm.n))*log(REAL(dpm.n))/2.0;
  os << "  sn =  " << fmt('f',20,8,sn) << fmt('e',27,17,sn) << '\n';
  os << "  aic = " << fmt('f',20,8,aic) << fmt('e',27,17,aic) << '\n';
  os << "  hq =  " << fmt('f',20,8,hq) << fmt('e',27,17,hq) << '\n';
  os << "  bic = " << fmt('f',20,8,bic) << fmt('e',27,17,bic) << '\n';
  
  realmat V = invpsd(infmat);
  
  realmat se(theta.size(),1,0.0);
  realmat tv(theta.size(),1);

  vector< pair<INTEGER,INTEGER> >::const_iterator rt_ptr; 
  for (rt_ptr = rt.begin(); rt_ptr != rt.end(); ++rt_ptr) {
    se[rt_ptr->second] = sqrt(V(rt_ptr->first,rt_ptr->first));
  }
  for (INTEGER i=1; i<=theta.size(); ++i) {
    tv[i] = (se[i]==0.0 ? 0.0 : theta[i]/se[i]); 
  }

  os << "Index           theta       std error     t-statistic    descriptor\n";
  INTEGER count=0;
  for (INTEGER i=1; i<=af.get_nrowa0(); ++i) {
    ++count;
    os << fmt('d',5,count) << ' '
       << fmt('f',15,5,theta[count]) << ' '
       << fmt('f',15,5,se[count]) << ' '
       << fmt('f',15,5,tv[count])
       << "    a0["<<i<<"]   ";
    for(INTEGER k=1; k<=f.get_ly(); ++k) {
      os << af.get_alpha()[af.get_odx()[i]-1][k];
    }
    os << '\n';
  }
  for (INTEGER j=1; j<=af.get_ncolA(); ++j) {
    for (INTEGER i=1; i<=af.get_nrowA(); ++i) {
      ++count;
      os << fmt('d',5,count) << ' '
         << fmt('f',15,5,theta[count]) << ' ' 
         << fmt('f',15,5,se[count]) << ' '
         << fmt('f',15,5,tv[count])
         << "    A("<<i<<','<<j<<")  ";
      for(INTEGER k=1; k<=f.get_ly(); ++k) {
        os << af.get_alpha()[af.get_idx()[i]-1][k];
      }
      os << " ";
      for(INTEGER k=1; k<=af.get_lx()*af.get_lag(); ++k) {
        os << af.get_beta()[j-1][k];
      }
      os << '\n';
    }
  }
  if (uf.is_intercept()) {
    for (INTEGER i=1; i<=uf.get_ly(); ++i) {
      ++count;
      os << fmt('d',5,count) << ' '
         << fmt('f',15,5,theta[count]) << ' ' 
         << fmt('f',15,5,se[count]) << ' '
         << fmt('f',15,5,tv[count])
         << "    b0["<<i<<"] ";
      os << '\n';
    }
  }
  if (uf.is_regression()) {
    for (INTEGER j=1; j<=uf.get_lx()*uf.get_lag(); ++j) {
      for (INTEGER i=1; i<=uf.get_ly(); ++i) {
        ++count;
        os << fmt('d',5,count) << ' '
           << fmt('f',15,5,theta[count]) << ' ' 
           << fmt('f',15,5,se[count]) << ' '
           << fmt('f',15,5,tv[count])
           << "    B("<<i<<','<<j<<") ";
        os << '\n';
      }
    }
  }
  for (INTEGER i=1; i<=rf.get_lR(); ++i) {
    ++count;
    os << fmt('d',5,count) << ' '
       << fmt('f',15,5,theta[count]) << ' ' 
       << fmt('f',15,5,se[count]) << ' '
       << fmt('f',15,5,tv[count])
       << "    R0["<<i<<"] ";
    os << '\n';
  }
  for (INTEGER j=1; j<=rf.get_colsP(); ++j) {
    for (INTEGER i=1; i<=rf.get_rowsP(); ++i) {
      ++count;
      os << fmt('d',5,count) << ' '
         << fmt('f',15,5,theta[count]) << ' ' 
         << fmt('f',15,5,se[count]) << ' '
         << fmt('f',15,5,tv[count])
         << "    P("<<i<<','<<j<<")  " << rf.get_Ptype();
      os << '\n';
    }
  }
  for (INTEGER j=1; j<=rf.get_colsQ(); ++j) {
    for (INTEGER i=1; i<=rf.get_rowsQ(); ++i) {
      ++count;
      os << fmt('d',5,count) << ' '
         << fmt('f',15,5,theta[count]) << ' ' 
         << fmt('f',15,5,se[count]) << ' '
         << fmt('f',15,5,tv[count])
         << "    Q("<<i<<','<<j<<")  " << rf.get_Qtype();
      os << '\n';
    }
  }
  for (INTEGER j=1; j<=rf.get_colsV(); ++j) {
    for (INTEGER i=1; i<=rf.get_rowsV(); ++i) {
      ++count;
      os << fmt('d',5,count) << ' '
         << fmt('f',15,5,theta[count]) << ' ' 
         << fmt('f',15,5,se[count]) << ' '
         << fmt('f',15,5,tv[count])
         << "    V("<<i<<','<<j<<")  " << rf.get_Vtype();
      os << '\n';
    }
  }
  for (INTEGER j=1; j<=rf.get_colsW(); ++j) {
    for (INTEGER i=1; i<=rf.get_rowsW(); ++i) {
      ++count;
      os << fmt('d',5,count) << ' '
         << fmt('f',15,5,theta[count]) << ' ' 
         << fmt('f',15,5,se[count]) << ' '
         << fmt('f',15,5,tv[count])
         << "    W("<<i<<','<<j<<")  " << rf.get_Wtype();
      os << '\n';
    }
  }

  os << '\n';
  os << "For estimation of this SNP specification by MCMC, "
     << "restrict the following\nelements of theta to be positive: ";
  for (INTEGER i=1; i<=srvec.size(); ++i) os << " " << srvec[i];
  os << '\n';
  os << '\n';

  os << "Largest eigen value of mean function companion matrix = " 
     << uf.stability() << '\n';

  os << "Largest eigen value of variance function P & Q companion matrix = " 
     << rf.stability() << '\n';
  os << '\n';

  os.flush();

  return os;
}

REAL snp::snpll::log_likehood() 
{
  if (rho.size()<=0) error("Error, snpll, can't use if default construced");
  if (Y==0||X==0) error("Error, snpll, log_likehood, data not initialized");
  if (dpm.M != Y->nrow() || dpm.M != X->nrow() || dpm.M != f.get_ly())
    error("Error, snpll, this should never happen");

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

  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;

    u = uf(x_lag);
    
    f.set_R(rf(y_lag,u_lag,x_lag));
    f.set_a(af(x_lag));
    f.set_u(u);
  }

  REAL log_likelihood = 0.0;

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

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

    u_lag = u;

    u = uf(x_lag);

    f.set_R(rf(y_lag,u_lag,x_lag));
    f.set_a(af(x_lag));
    f.set_u(u);

    log_likelihood += f.log_f(y);
  }

  if (compute_infmat) {
    error ("Error, snpll, method likehood() cannot compute infmat");
  }

  return log_likelihood;
}

void snp::snpll::rho_hessian(scl::realmat& rho_hessian)
{
  realmat rho_save = get_rho();
  realmat x = rho_save;
  INTEGER d = x.size();
  rho_hessian.resize(d,d);
  realmat f0;
  realmat f1;
  REAL eps = std::pow(REAL_EPSILON,0.33333333);
  for (INTEGER j=1; j<=d; j++) {
    REAL tmp = x[j];
    REAL h = eps*std::fabs(tmp);
    if (h == 0) h = eps;
    REAL hi = tmp + h;
    REAL lo = tmp - h;
    x[j] = hi;
    set_rho(x);
    log_likehood(f1);
    x[j] = lo;
    set_rho(x);
    log_likehood(f0);
    REAL difference = hi - lo;
    x[j] = tmp;
    for (INTEGER i=1; i<=d; i++) {
      rho_hessian(i,j) = (f1[i] - f0[i])/difference;
    }
  }
  set_rho(rho_save);
}
