#include "libscl.h"
#include "tree_sim.h"
#include "tree_mf.h"
#include "preimage_norm.h"

using namespace std;
using namespace scl;

extern void rom(INTEGER rows, INTEGER cols, scl::realmat& Q, INT_32BIT& seed);

int main(int argc, char** argp, char** envp)
{
  const INTEGER n_obs = 100;
  const INTEGER dat_rows = 2;
  const INTEGER mf_lags = 1;
  const INTEGER HAC_lags = 1;
  const INTEGER n_prior_reps = 500;

  INT_32BIT seed = 740726;

  realmat rho(4,1);
  realmat prior_mean(4,1);
  realmat prior_sdev(4,1);
  realmat theta(2,1);

  REAL alpha = prior_mean[1] = 0.95;
  REAL sigma = prior_mean[2] = 0.02;
  REAL beta  = theta[1] = prior_mean[3] = 0.95;
  REAL gamma = theta[2] = prior_mean[4] = 12.5;

  prior_sdev[1] = 0.01;
  prior_sdev[2] = 0.01;
  prior_sdev[3] = 0.01;
  prior_sdev[4] = 2.00;

  INTEGER scope = 16;
  REAL increment = 0.02;
  REAL tol = 1.0e-3;
  REAL iter_limit = 5000;

  cout << '\n';
  cout << "Parameters" << '\n';
  cout << "alpha_mean " << fmt('e',26,16,alpha) << "   " << alpha << '\n';
  cout << "sigma_mean " << fmt('e',26,16,sigma) << "   " << sigma << '\n';
  cout << "beta_mean  " << fmt('e',26,16,beta)  << "   " << beta  << '\n';
  cout << "gamma_mean " << fmt('e',26,16,gamma) << "   " << gamma << '\n';
  cout << "alpha_sdev " << fmt('e',26,16,prior_sdev[1])
       << "   " << prior_sdev[1] << '\n';
  cout << "sigma_sdev " << fmt('e',26,16,prior_sdev[2])
       << "   " << prior_sdev[2] << '\n';
  cout << "beta_sdev  " << fmt('e',26,16,prior_sdev[3])
       << "   " << prior_sdev[3] << '\n';
  cout << "gamma_sdev " << fmt('e',26,16,prior_sdev[4])
       << "   " << prior_sdev[4] << '\n';
  cout << '\n';
  cout << "Sizes" << '\n';
  cout << "n_obs " << n_obs << '\n';
  cout << "mf_lags " << mf_lags << '\n';
  cout << "HAC_lags " << HAC_lags << '\n';
  cout << "n_prior_reps " << n_prior_reps << '\n';
  cout << "scope " << scope << '\n';
  cout << "increment " << increment << '\n';
  cout << "tol " << tol << '\n';
  cout << "iter_limit " << iter_limit << '\n';
  cout << "scope " << scope << '\n';
  cout << '\n';

  bool rv = true;

  realmat data(dat_rows,n_obs);
  realmat data_orig(dat_rows,n_obs);
  realmat data_tran(n_obs,dat_rows);

  tree_variables mv;

  bool success = tree_sim(prior_mean, n_obs, seed, mv);

  if (!success) rv = false;

  for (INTEGER t=1; t<=n_obs; ++t) {
    data(1,t) = data_orig(1,t) = data_tran(t,1) = mv.geometric_stock_return[t];
    data(2,t) = data_orig(2,t) = data_tran(t,2) = mv.log_consumption_growth[t];
  }

  intvec idx = seq(1,2);
  realmat b0,B,V,C;

  //varcoef(data_tran,idx,1,b0,B,V,C);
  //cout << B;

  //cout << simple(T(data));

  preimage_norm Znorm(data,n_obs,mf_lags,HAC_lags);

  INTEGER d = Znorm.get_d();

  realmat z(d,1);
  realmat f(1,1);

  realmat x(scope,1);
  realmat xtst(scope,1);
  realmat xmin(scope,1);

  tree_moment_function tree_mf;

  gmm tree_gmm(&tree_mf, mf_lags, &data, n_obs, HAC_lags);

  tree_gmm.set_correct_W_for_mean(true);
  tree_gmm.set_warning_messages(true);

  realmat m;
  realmat W;
  REAL logdetW;
  INTEGER rankW;
  realmat S;

  realmat m_sum(d,1,0.0);
  realmat m_ssq(d,d,0.0);

  for (INTEGER rep=1; rep<=n_prior_reps; ++rep) {

    for (INTEGER i=1; i<=4; ++i) {
      rho[i] = prior_mean[i] + prior_sdev[i]*unsk(seed);
    }

    while (fabs(rho[1]) >= 1.0) rho[1]=prior_mean[1]+prior_sdev[1]*unsk(seed);
    while (rho[2] <= 0.0) rho[2]=prior_mean[2]+prior_sdev[2]*unsk(seed);
    while ((rho[3] <= 0.0) || (rho[3] >= 1.0))
      rho[3]=prior_mean[3]+prior_sdev[3]*unsk(seed);
    while (rho[4] <= 0.0) rho[4]=prior_mean[4]+prior_sdev[4]*unsk(seed);

    beta  = theta[1] = prior_mean[3];
    gamma = theta[2] = prior_mean[4];

    for (INTEGER i=1; i<=d; ++i) {
      z[i] = unsk(seed);
    }

    Znorm.set_z_theta(z,theta);
 
    data = data_orig;

    Znorm.set_data(data);

    INTEGER t = n_obs/2 - scope;
    INTEGER base = dat_rows*(t - 1);
    for (INTEGER i=1; i<=scope; ++i) x[i] = data[base + i];

    Znorm.set_t(t);

    REAL ymin;
    if (!Znorm.get_f(x,f)) rv = false;
    ymin = f[1]; 
    xmin = x;

    INTEGER iter = 1;
    cout << '\n';
    cout << fmt('d',5,iter) <<' '<< fmt('f',20,8,ymin) << '\n';
    while (iter<=iter_limit) {

      ++iter;

      ++t;

      base = dat_rows*(t - 1);

      if (base + scope  > n_obs*dat_rows) t = 1;

      base = dat_rows*(t - 1);

      Znorm.set_t(t);

      for (INTEGER i=1; i<=scope; ++i) x[i] = data[base + i];

      realmat Q;
      rom(scope,scope,Q,seed);

      REAL incr = ran(seed)*increment;

      for (INTEGER j=1; j<=scope; ++j) {
        for (INTEGER i=1; i<=scope; ++i) {
          xtst[i] = x[i] + incr*Q(i,j);
        }
        if (Znorm.get_f(xtst,f)) { 
          if (f[1] < ymin) {ymin = f[1]; xmin = xtst;}
        }
        else {
          Znorm.set_x(x);
        }
      }

      Znorm.set_x(xmin); 

      data = Znorm.get_data();

      if (ymin < tol) break;
    }

    cout << fmt('d',5,iter) <<' '<< fmt('f',20,8,ymin) << '\n';
    cout.flush();

    for (INTEGER t=1; t<=n_obs; ++t) {
      for (INTEGER i=1; i<=dat_rows; ++i) {
        data_tran(t,i) = data(i,t);
      }
    }

    //varcoef(data_tran,idx,1,b0,B,V,C);
    //cout << B;

    //cout << simple(T(data));

    if (iter >= iter_limit) rv = false;

    if(!tree_gmm.set_data(&data)) rv = false;

    tree_gmm(prior_mean,m,W,logdetW,rankW,S);

    m_sum += m;
    m_ssq += m*T(m);
  }

  realmat W_pre = m_ssq - m_sum*T(m_sum)/n_prior_reps;

  W_pre = W_pre/(n_prior_reps - 1);

  vecwrite("W_pre.dat", W_pre);

  cout << (rv ? "success" : "failure") << '\n';

  return 0;
}

