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

Copyright (C) 2013.

A. Ronald Gallant

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 "libmle.h"
using namespace std;  
using namespace scl;
using namespace libmle;

#define ADJUST_FOR_PRIOR

// Conceptually the prior is treated as the density of the first observation.

bool libmle::minimal_asymptotics::set_asymptotics(const realmat& chain)
{
  INTEGER n_chain = chain.ncol();

  if (chain.nrow()!=p) error("Error, minimal_asymptotics, bad chain matrix");

  if (n_tot == 0) {
    for (INTEGER t=1; t<=n_chain; ++t) {
      for (INTEGER i=1; i<=p; i++) {
        rho_ctr[i] += chain(i,t);
      }
    }
    rho_ctr = rho_ctr/n_chain;
  }

  for (INTEGER t=1; t<=n_chain; ++t) {
    for (INTEGER i=1; i<=p; i++) {
      rho_sum[i] += (chain(i,t) - rho_ctr[i]);
    }
  }

  for (INTEGER t=1; t<=n_chain; ++t) {
    for (INTEGER j=1; j<=p; j++) {
      for (INTEGER i=1; i<=p; i++) {
        rho_sse(i,j) += (chain(i,t) - rho_ctr[i])*(chain(j,t) - rho_ctr[j]);
      }
    }
  }

  n_tot += n_chain;

  mean = rho_sum/n_tot; 
  sscp = rho_sse/n_tot;

  for (INTEGER j=1; j<=p; j++) {
    for (INTEGER i=1; i<=p; i++) {
      sscp(i,j) -= mean[i]*mean[j];
    }
  }
  mean += rho_ctr;
  sscp = REAL(n_dat)*mcmc.get_temp()*sscp;

  realmat mode_new = mcmc.get_mode();
  REAL high_new = mcmc.get_high();
  if (high_new > high) {
    mode = mode_new;
    high = high_new;;
  }

  return true;
}

void minimal_asymptotics::get_asymptotics (realmat& rho_hat, realmat& V_hat,
         INTEGER& n)
{ 
  rho_hat = mean;
  V_hat = sscp/n_dat;
  n = n_dat;
}

void minimal_asymptotics::get_asymptotics(realmat& rho_mean, realmat& rho_mode,
         REAL& post_high, realmat& I, realmat& invJ, realmat& foc_hat, 
         INTEGER& reps)
{ 
  realmat null;
  rho_mean = mean; 
  rho_mode = mode; 
  post_high = high;
  I = null; 
  invJ = sscp;
  foc_hat = null;
  reps = 0;
}

namespace {

  REAL parzen(REAL x)
  {
    REAL z = fabs(x);
    if (z <= 0.5) {
      return 1.0 - 6.0*pow(z,2) + 6.0*pow(z,3);
    }
    if (z <= 1.0) {
      return 2.0*pow((1.0 - z),3);
    }
    return 0.0;
  }

}

bool libmle::sandwich_asymptotics::set_asymptotics(const realmat& chain)
{
  INTEGER n_chain = chain.ncol();

  if (chain.nrow()!=p) error("Error, sandwich_asymptotics, bad chain matrix");

  if (n_tot == 0) {
    for (INTEGER t=1; t<=n_chain; ++t) {
      for (INTEGER i=1; i<=p; i++) {
        rho_ctr[i] += chain(i,t);
      }
    }
    rho_ctr = rho_ctr/n_chain;
  }

  for (INTEGER t=1; t<=n_chain; ++t) {
    for (INTEGER i=1; i<=p; i++) {
      rho_sum[i] += (chain(i,t) - rho_ctr[i]);
    }
  }

  for (INTEGER t=1; t<=n_chain; ++t) {
    for (INTEGER j=1; j<=p; j++) {
      for (INTEGER i=1; i<=p; i++) {
        rho_sse(i,j) += (chain(i,t) - rho_ctr[i])*(chain(j,t) - rho_ctr[j]);
      }
    }
  }

  n_tot += n_chain;

  mean = rho_sum/n_tot; 
  sscp = rho_sse/n_tot;

  for (INTEGER j=1; j<=p; j++) {
    for (INTEGER i=1; i<=p; i++) {
      sscp(i,j) -= mean[i]*mean[j];
    }
  }
  mean += rho_ctr;
  sscp = REAL(n_dat)*mcmc.get_temp()*sscp;

  realmat mode_new = mcmc.get_mode();
  REAL high_new = mcmc.get_high();
  if (high_new > high) {
    mode = mode_new;
    high = high_new;;
  }

  model.set_rho(mode);
  
  realmat scores;

  if (model.get_scores(scores)) {

    foc.resize(p,1,0.0);
    for (INTEGER t=1; t<=n_dat; ++t) {
      for (INTEGER i=1; i<=p; ++i) {
        foc[i] += scores(i,t);
      }
    }

    realmat R0(p,p,0.0);
    for (INTEGER t=1; t<=n_dat; ++t) {
      for (INTEGER j=1; j<=p; ++j) {
        for (INTEGER i=1; i<=p; ++i) {
          R0(i,j) += scores(i,t)*scores(j,t);
        }
      }
    }

    I_mat = R0;

    realmat R1(p,p);

    for (INTEGER lag=1; lag<=lhac; lag++) {
    
      fill(R1);

      for (INTEGER t=1+lag; t<=n_dat; t++) {
        for (INTEGER j=1; j<=p; ++j) {
          for (INTEGER i=1; i<=p; ++i) {
            R1(i,j) += scores(i,t)*scores(j,t-lag);
          }
        }
      }

      I_mat += parzen(REAL(lag)/REAL(lhac))*(R1 + T(R1));
    }

    I_reps = n_dat;

    return true;
  }
  else {
   warn("Warning, sandwich_asymptotics, cannot compute scores");
   realmat null;
   I_mat = null;
   foc = null;
   I_reps = 0;
   return false;
 }
}

void libmle::sandwich_asymptotics::get_asymptotics (realmat& rho_hat, 
         realmat& V_hat, INTEGER& n)
{ 
  rho_hat = mode;
  realmat invJ = sscp;
  if (I_mat.size() > 0) {
    realmat I = I_mat/I_reps;
    V_hat = (invJ*I*invJ)/n_dat;
  }
  else {
    V_hat = sscp/n_dat;
  }
  n = n_dat;
}

void libmle::sandwich_asymptotics::get_asymptotics (realmat& rho_mean, 
         realmat& rho_mode, REAL& post_high, realmat& I, realmat& invJ, 
         realmat& foc_hat, INTEGER& reps)
{ 
  rho_mean = mean; 
  rho_mode = mode; 
  post_high = high;
  if (I_mat.size() > 0) {
    I = I_mat/I_reps; 
  }
  else {
    I = I_mat;
  }
  invJ = sscp;
  foc_hat = foc;
  reps = I_reps;
}

