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

Copyright (C) 2009

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 <cerrno>
#include "libsmm.h"
#include "emm.h"

using namespace scl;
using namespace libsmm;
using namespace emm;
using namespace std;

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

  void 
  xvec(const realmat& dat, const realmat& mu, INTEGER L, INTEGER t, realmat& x)
  {
    if (t <= L) error("Error, xvec, bad index");

    INTEGER d = dat.nrow();
    INTEGER len = d + (d*d + d)/2 + L*(d*d);
    x.resize(len,1);

    INTEGER off = 0;

    for (INTEGER i = 1; i <= d; i++) {
      x[i] = dat(i,t);
    }

    off += d;

    for (INTEGER j=1; j<=d; j++) {
      for (INTEGER i=1; i<=j; i++) {
        x[off + (j*j-j)/2 + i] = (dat(i,t) - mu[i]) * (dat(j,t) - mu[j]);
      }
    }

    off += (d*d + d)/2;

    for (INTEGER k = 1; k <= L; ++k) {

      for (INTEGER j=1; j<=d; j++) {
        for (INTEGER i=1; i<=d; i++) {
          x[off + d*(j-1) + i] = (dat(i,t) - mu[i]) * (dat(j,t-k) - mu[j]);
        }
      }

      off += d*d;
    }
  }

} 

namespace emm {

  gmm_objfun::gmm_objfun
    (const realmat& data, INTEGER num_mod_parms, INTEGER num_mod_funcs,
     const std::vector<std::string>& mod_pfvec,
     const std::vector<std::string>& mod_alvec,
     const std::vector<std::string>& obj_pfvec,
     const std::vector<std::string>& obj_alvec,
     std::ostream& detail)
  {
    if (obj_alvec.size() != 5) {
      error("Error, gmm_objfun, L, Lhac, or obj_print not in parmfile"); 
    }

    L = atoi(obj_alvec[1].substr(0,12).c_str());
    Lhac = atoi(obj_alvec[2].substr(0,12).c_str());
    bool obj_print = atoi(obj_alvec[3].substr(0,12).c_str());
    
    set_data(data);

    if (obj_print) {
      detail << starbox("/Begin gmm_objfun constructor output//");
      detail << '\n';
      detail << "\t d = " << d << '\n';
      detail << "\t n = " << n << '\n';
      detail << "\t L = " << L << '\n';
      detail << "\t Lhac = " << Lhac << '\n';
      detail << '\n';
      detail << "\t m = " << m;
      detail << "\t W = " << W;
      detail << starbox("/First 12 observations//");
      detail << data("",seq(1,12));
      detail << starbox("/Last 12 observations//");
      detail << data("",seq(n-11,n));
      detail << starbox("/End gmm_objfun constructor output//");
      detail << '\n';
      detail.flush();
    }
  }

  void gmm_objfun::set_data(const realmat& dat)
  {
    d = dat.nrow();
    n = dat.ncol();
  
    realmat mu(d,1,0.0);
  
    for (INTEGER t=1; t<=n; t++) {
      for (INTEGER i=1; i<=d; i++) {
        mu[i] += dat(i,t);
      }
    }
  
    mu = mu/n;
  
    realmat x;
    xvec(dat,mu,L,L+1,x);
  
    INTEGER len_m = x.size();
  
    m.resize(len_m,1,0.0);
  
    for (INTEGER t=L+1; t<=n; t++) {
      xvec(dat,mu,L,t,x);
      m += x;
    }
  
    m = m/n;
  
    if (L >= n-1) {
      warn("Warning, L too large, set to n-2 instead");
      L = (n-2 > 0) ? n-2 : 0;
    }
  
    realmat R0(len_m,len_m,0.0);
  
    for (INTEGER t=L+1; t<=n; t++) {
      xvec(dat,mu,L,t,x);
      R0 += (x-m)*T(x-m);
    }
  
    R0 = R0/n;
  
    W = R0;
    realmat R1(len_m,len_m);
  
    for (INTEGER lag=1; lag<=Lhac; lag++) {
    
      fill(R1);
      
      realmat xlag;
      for (INTEGER t=L+1+lag; t<=n; t++) {
        xvec(dat,mu,L,t,x);
        xvec(dat,mu,L,t-lag,xlag);
        R1 += (x-m)*T(xlag-m);
      }
  
      R1 = R1/n;
  
      W += parzen(REAL(lag)/REAL(L))*(R1 + T(R1)); 
    }
  
    W = invpsd(W);
  }
  
  REAL gmm_objfun::operator()(const realmat& rho, const realmat& sim) const
  {
    if (sim.nrow() != d) error("Error, gmm_objfun, bad dimension");
  
    INTEGER len_sim = sim.ncol();
  
    realmat mu_sim(d,1,0.0);
  
    for (INTEGER t=1; t<=len_sim; t++) {
      for (INTEGER i=1; i<=d; i++) {
        mu_sim[i] += sim(i,t);
      }
    }
  
    mu_sim = mu_sim/len_sim;
  
    realmat x;
    xvec(sim,mu_sim,L,L+1,x);
  
    INTEGER len_m = x.size();
  
    realmat m_sim(len_m,1,0.0);
  
    for (INTEGER t=L+1; t<=len_sim; t++) {
      xvec(sim,mu_sim,L,t,x);
      m_sim += x;
    }
  
    m_sim = m_sim/len_sim;
  
    realmat r = (T(m - m_sim)*W)*(m - m_sim);
  
    REAL v = (REAL(n)/2.0)*r[1];
  
    return v;
  }  

  objfun_base* gmm_objfun::new_objfun() 
  {
     objfun_base* r = new(nothrow) gmm_objfun(*this);
     if (r == 0) error("Error, gmm_objfun, operator new failed");
     return r;
  }

  var_usrmod::var_usrmod(const scl::realmat& dat, 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)
  : b(2,1), B(2,2), R(2,2), 
    n_stats(2), n_parms(2+4+3), n_datum(2), bseed(740726),
    statistics(2,1,0.0)
  {
    if (len_mod_parm != n_parms) error("Error, var_usrmod, bad len_mod_parm");
    if (len_mod_func != n_stats) error("Error, var_usrmod, bad len_mod_parm");
    if (dat.nrow() != n_datum) error("Error, var_usrmod, bad M");

    n_obser = dat.ncol();

    b[1] = 0.1; b[2] = 0.1;
    B[1] = 0.5; B[2] = 0.2; B[3] = 0.2; B[4] = 0.5;
    R[1] = 0.001; R[2] = 0.0; R[3] = 0.0001; R[4] = 0.001;

    sim_size = 50000;
    drop = 500;

    if (mod_alvec.size() != 5) {
      error("Error, var_usrmod, slen, spin, or mod_print not in parmfile");
    }

    sim_size = atoi(mod_alvec[1].substr(0,12).c_str());
    drop = atoi(mod_alvec[2].substr(0,12).c_str());
    bool mod_print = atoi(mod_alvec[3].substr(0,12).c_str());
    
    if (mod_print) {
      detail << starbox("/Begin var_usrmod constructor output//");
      detail << '\n';
      detail << "\t n_stats = " << n_stats << '\n';
      detail << "\t n_parms = " << n_parms << '\n';
      detail << "\t n_datum = " << n_datum << '\n';
      detail << "\t n_obser = " << n_obser << '\n';
      detail << "\t sim_size = " << sim_size << '\n';
      detail << "\t drop = " << drop << '\n';
      detail << starbox("/End var_usrmod constructor output//");
      detail << '\n';
      detail.flush();
    }
    
  }

  bool var_usrmod::make_sim(INT_32BIT& seed, realmat& sim) 
  {
    if (n_datum != 2) error("Error, var_usrmod, make_sim, n_datum wrong");
    
    sim.resize(n_datum,sim_size);
  
    realmat u(2,1);
    realmat y(2,1);
    realmat y_lag(2,1,0.0);
  
    for (INTEGER t=1; t<=drop; t++) {
      u[1] = unsk(&seed);
      u[2] = unsk(&seed);
      y = b + B*(y_lag-b) + R*u;
      y_lag = y;
    }
  
    for (INTEGER t=1; t<=sim_size; t++) {
      u[1] = unsk(&seed);
      u[2] = unsk(&seed);
      y = b + B*(y_lag-b) + R*u;
      y_lag = y;
      sim(1,t) = y[1];
      sim(2,t) = y[2];
    }
  
    return true;
  }
  
  bool var_usrmod::gen_sim(realmat& sim, realmat& stats)
  {
    stats = statistics;
    INT_32BIT seed = fixed_seed;
    return make_sim(seed,sim);
  }
  
  bool var_usrmod::gen_bootstrap(vector<realmat>& bs)
  {
    realmat sim;
    if (!make_sim(bseed,sim)) return false;;

    INTEGER possible = sim.ncol()/(n_obser+drop);
    INTEGER desired  = 2*n_parms+1;

    if (possible < desired) {
      error("Error, var_usrmod, gen_bootstrap, slen too small or spin too big");
    }
    
    vector<realmat>::size_type len_vec = desired;
    bs.resize(len_vec);

    realmat dat(n_datum,n_obser);
    INTEGER s = 0;
    for (vector<realmat>::size_type k=0; k<len_vec; ++k) {
      for (INTEGER j=1; j<=n_obser; ++j) {
        for (INTEGER i=1; i<=n_datum; ++i) {
	  dat(i,j) = sim(i,s+j);
        }
      }
      bs[k] = dat;
      s += n_obser + drop;
    }

    return true;
  }

  bool var_usrmod::support(const realmat& rho)
  {
    realmat E;
    INTEGER ier;
    return (eigen(B,E,ier) < 1.0) && (0.0 < R[1]) && (0.0 < R[4]);
  }
  
  den_val var_usrmod::prior(const realmat& theta, const realmat& stats)
  {
    /*
    const REAL minus_log_root_two_pi = -9.1893853320467278e-01;
    realmat Q = T(theta)*theta;
    return den_val(true, REAL(n_parms)*minus_log_root_two_pi - 0.5*Q[1]);
    */
    return den_val(true, 0.0);
  }

}