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

using namespace std;
using namespace scl;

int main(int argc, char** argp, char** envp)
{
  const INTEGER n_obs = 500;
  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);

  REAL alpha = prior_mean[1] = 0.95;
  REAL sigma = prior_mean[2] = 0.03;
  REAL beta  = prior_mean[3] = 0.95;
  REAL gamma = prior_mean[4] = 1.2600114143005971e+01;

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

  cout << '\n';
  cout << "Parameters" << '\n';
  cout << "alpha_mean " << fmt('e',27,17,alpha) << '\n';
  cout << "sigma_mean " << fmt('e',27,17,sigma) << '\n';
  cout << "beta_mean  " << fmt('e',27,17,beta)  << '\n';
  cout << "gamma_mean " << fmt('e',27,17,gamma) << '\n';
  cout << "alpha_sdev " << fmt('e',27,17,prior_sdev[1]) << '\n';
  cout << "sigma_sdev " << fmt('e',27,17,prior_sdev[2]) << '\n';
  cout << "beta_sdev  " << fmt('e',27,17,prior_sdev[3]) << '\n';
  cout << "gamma_sdev " << fmt('e',27,17,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';

  bool rv = true;

  realmat data(2,n_obs);
  realmat data_orig(2,n_obs);

  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) = mv.geometric_stock_return[t];
    data(2,t) = data_orig(2,t) = mv.log_consumption_growth[t];
  }

  preimage_eqns Zeqns(data,n_obs,mf_lags,HAC_lags);

  INTEGER d = Zeqns.get_d();

  INTEGER r = data.nrow();
  INTEGER last = r*n_obs;

  realmat x_start(d,1);
  realmat x_stop(d,1);

  for (INTEGER i=0; i<d; ++i) {
    x_start[d - i] = data[last - i];
  }

  nlsolve Zinv(Zeqns);

  realmat z(d,1);

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

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

    Zeqns.update_z_rho(z,rho);

    Zinv.solve(x_start, x_stop);
  }

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

  return 0;
}

