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

using namespace std;
using namespace scl;

realmat mfun(INTEGER t, const tree_variables& mv, const realmat& theta)
{
  INTEGER d = 3;
  INTEGER T0 = 4;
  if (t < T0) error("Error, mfun, t < T0");
  REAL beta = theta[1];
  REAL gamma = theta[2];
  REAL c_lag = mv.log_consumption[t-1];
  REAL c = mv.log_consumption[t];
  REAL p_lag = mv.log_stock_price[t-1];
  REAL p = mv.log_stock_price[t];
  REAL lmrs = log(beta) - gamma*(c - c_lag);
  REAL mrs = exp(lmrs);
  REAL payoff = exp(p) + exp(c);
  REAL price = exp(p_lag);
  REAL error = 1.0 - mrs*(payoff/price);
  realmat mt(d,1);
  mt[1] = error;
  mt[2] = p_lag*error;
  mt[3] = c_lag*error;
  return mt;
}

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

  INT_32BIT seed = 740726;

  realmat prior_mean(4,1);
  realmat prior_sdev(4,1);

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

  REAL lo =  -3.5*sigma/sqrt(1.0 - alpha*alpha);
  REAL hi = -lo;

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

  cout << '\n';
  cout << "Prior mean" << '\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 << '\n';
  cout << "Prior sdev" << '\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';

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

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

  bool rv = true;

  tree_variables mv;

  bool success = tree_sim(prior_mean, n_obs, seed, mv);
  if (!success) rv = false;

  realmat c = mv.log_consumption;
  REAL min_c = REAL_MAX;
  REAL max_c = -REAL_MAX;
  for (INTEGER i=1; i<=n_obs; ++i) {
    min_c = c[i] < min_c ? c[i] : min_c;
    max_c = c[i] > max_c ? c[i] : max_c;
  }

  cout << '\n';
  cout << "Policy function domain" << '\n';
  cout << "lo " << fmt('e',26,16,lo) << "   " << lo << '\n';
  cout << "hi " << fmt('e',26,16,hi) << "   " << hi << '\n';
  cout << "log_consumption range" << '\n';
  cout << "min_c " << fmt('e',26,16,min_c) << "   " << min_c << '\n';
  cout << "max_c " << fmt('e',26,16,max_c) << "   " << max_c << '\n';

  realmat tree_vars;
  vector<string> mv_names = mv.get_tree_variables(tree_vars);
  writetable("tree_vars.txt",tree_vars,mv_names,35,15);

  INTEGER T0 = 4;
  INTEGER  d = 3;

  realmat m(d,1);
  realmat S(d,d);
  realmat R;
  realmat Rinv;
  realmat z(d,1);

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

  fill(m);
  for (INTEGER t=T0; t<=n_obs; ++t) {
    m += mfun(t,mv,theta);
   }
   m = m/(n_obs - T0);

  realmat mvals(n_obs,d,0.0);
  realmat maxvals(d,1,-REAL_MAX);
  realmat minvals(d,1,REAL_MAX);

  fill(S);
  realmat mt(d,1);
  for (INTEGER t=T0; t<=n_obs; ++t) {
    mt = mfun(t,mv,theta);
    for (INTEGER i=1; i<=d; ++i) {
      mvals(t,i) = mt[i];
      maxvals[i] = maxvals[i] < mt[i] ? mt[i] : maxvals[i];
      minvals[i] = minvals[i] > mt[i] ? mt[i] : minvals[i];
    }
    S += mt*T(mt);
  }
  S = S/(n_obs - T0);

  if (cholesky(S,R) != d) rv = false;
  rinv(R,Rinv);
  z = sqrt(n_obs)*T(Rinv)*m;

  vector<string> header(d);
  header[0] = "m[1]";
  header[1] = "m[2]";
  header[2] = "m[3]";

  writetable("mvals.txt",mvals,header,20,8);

  cout << '\n';
  cout << "First z computation" << '\n';
  cout << "m " << m << '\n';
  cout << "S " << S << '\n';
  cout << "R " << R << '\n';
  cout << "Rinv " << Rinv << '\n';
  cout << "z " << z << '\n';
  cout << '\n';

  realmat stats = simple(mvals);

  cout << "range " << cbind(minvals,maxvals) << '\n';
  cout << "simple stats " << stats << '\n';;


/*
  intvec idx = seq(T0,n_obs);
  realmat stats = simple(mv.pricing_errors(idx,""));
  //realmat stats = simple(mv.pricing_errors);
  REAL pesum = 0.0;
  for (INTEGER i=T0; i<=mv.pricing_errors.size(); ++i) {
    pesum += mv.pricing_errors[i];
  }
  cout << "pricing_errors mean  " << pesum/mv.pricing_errors.size() << '\n';
  cout << "pricing_errors mean  " 
       << REAL(n_obs-T0+1)/REAL(n_obs)*stats[1] << '\n';
  cout << "pricing_errors sdev  " << stats[2] << '\n';
  cout << "pricing_errors var   " << stats[3] << '\n';
  cout << "pricing_errors skew  " << stats[4] << '\n';
  cout << "pricing_errors kurt  " << stats[5] << '\n';
  REAL sum = 0.0;
  for (INTEGER i=T0; i<=n_obs; ++i) {
    sum += (mv.pricing_errors[i]-stats[1])*(mv.pricing_errors[i-1]-stats[1]);
  }
  sum /= n_obs;
  cout << "pricing_errors auto  " << sum/stats[3] << '\n';
  cout << '\n';

  realmat S_cu(n_prior_reps,d*d);
  realmat W_cu(n_prior_reps,d*d);
  realmat Z_cu(n_prior_reps,d);

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

  cout << "Statistics over " << n_prior_reps << " repetitions";
  cout << '\n';

  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]) >= 0.99) rho[1]=prior_mean[1]+prior_sdev[1]*unsk(seed);
    while ((rho[2] <= 0.01) || (rho[2] >= 100.0)) 
      rho[2]=prior_mean[2]+prior_sdev[2]*unsk(seed);
    while ((rho[3] <= 0.80) || (rho[3] >= 0.99)) 
      rho[3]=prior_mean[3]+prior_sdev[3]*unsk(seed);
    while ((rho[4] <= 0.01) || (rho[4] >= 100.0))
      rho[4]=prior_mean[4]+prior_sdev[4]*unsk(seed);
    
    alpha = rho[1];
    sigma = rho[2];
    beta  = theta[1] = rho[3];
    gamma = theta[2] = rho[4];

    success = tree_sim(rho, n_obs, seed, mv);
    if (!success) rv = false;

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

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

    obj = tree_gmm(theta,m,W,logdetW,rankW,S);

    for (INTEGER j=1; j<=d*d; ++j) {
      S_cu(rep,j) = S[j];
      W_cu(rep,j) = W[j];
    }
    
    m_sum += m;
    m_ssq += m*T(m);

    if (cholesky(S,R) != d) rv = false;
    rinv(R,Rinv);
    realmat z = sqrt(n_obs)*T(Rinv)*m;

    for (INTEGER j=1; j<=d; ++j) {
      Z_cu(rep,j) = z[j];
    } 

    if (Z_cu(rep,2) > 3000) {
      cout << "Large deviation" << '\n';
      cout << rho;
      cout << m;
      cout << S;
      cout << W;
      cout << mv.pricing_errors;
    }

  }

  if (!rv) warn("Warning, rv is false");

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

  S_m = S_m/(n_prior_reps - 1);

  vecwrite("S_m.dat", S_m);

  realmat S_cu_stats =  simple(S_cu);

  realmat S_cu_mean(d,d);
  realmat S_cu_sdev(d,d);
  realmat S_cu_skew(d,d);
  realmat S_cu_kurt(d,d);

  for (INTEGER i=1; i<=d*d; ++i) {
    S_cu_mean[i] = S_cu_stats(1,i);
    S_cu_sdev[i] = S_cu_stats(2,i);
    S_cu_skew[i] = S_cu_stats(3,i);
    S_cu_kurt[i] = S_cu_stats(4,i);
  }

  vecwrite("S_cu_mean.dat", S_cu_mean);
  vecwrite("S_cu_sdev.dat", S_cu_sdev);

  realmat W_cu_stats =  simple(W_cu);

  realmat W_cu_mean(d,d);
  realmat W_cu_sdev(d,d);
  realmat W_cu_skew(d,d);
  realmat W_cu_kurt(d,d);

  for (INTEGER i=1; i<=d*d; ++i) {
    W_cu_mean[i] = W_cu_stats(1,i);
    W_cu_sdev[i] = W_cu_stats(2,i);
    W_cu_skew[i] = W_cu_stats(3,i);
    W_cu_kurt[i] = W_cu_stats(4,i);
  }

  cout << '\n';
  cout << "Weighting matrix" << '\n';
  cout << "S_m " << S_m;
  cout << "Mean of S_cu " << S_cu_mean << '\n';
  cout << "Sdev of S_cu " << S_cu_sdev << '\n';
  cout << "Skew of S_cu " << S_cu_skew << '\n';
  cout << "Kurt of S_cu " << S_cu_kurt << '\n';
  cout << "Mean of W_cu " << W_cu_mean << '\n';
  cout << "Sdev of W_cu " << W_cu_sdev << '\n';
  cout << "Skew of W_cu " << W_cu_skew << '\n';
  cout << "Kurt of W_cu " << W_cu_kurt << '\n';

  realmat Z_cu_stats =  simple(Z_cu);

  realmat Z_cu_mean(d,1);
  realmat Z_cu_sdev(d,1);
  realmat Z_cu_skew(d,1);
  realmat Z_cu_kurt(d,1);

  for (INTEGER i=1; i<=d; ++i) {
    Z_cu_mean[i] = Z_cu_stats(1,i);
    Z_cu_sdev[i] = Z_cu_stats(2,i);
    Z_cu_skew[i] = Z_cu_stats(4,i);
    Z_cu_kurt[i] = Z_cu_stats(5,i);
  }

  cout << '\n';
  cout << "Z random variable" << '\n';
  cout << "Mean of m    " << m_sum/n_prior_reps;
  cout << "Mean of Z_cu " << Z_cu_mean << '\n';
  cout << "Sdev of Z_cu " << Z_cu_sdev << '\n';
  cout << "Skew of Z_cu " << Z_cu_skew << '\n';
  cout << "Kurt of Z_cu " << Z_cu_kurt << '\n';

  vector<string> header(d);
  header[0] = "e";
  header[1] = "lsrlag*e";
  header[2] = "lcglag*e";

  writetable("Z.txt",Z_cu,header,20,8);

  realmat s = simple(mv.log_consumption_growth);

  cout << '\n';
  cout << "log_consumption_growth" << '\n';
  cout << "mean = " << fmt('f',10,5,s[1]) << '\n';
  cout << "sdev = " << fmt('f',10,5,s[2]) << '\n';

  s = simple(mv.geometric_stock_return);

  cout << '\n';
  cout << "geometric_stock_return" << '\n';
  cout << "mean = " << fmt('f',10,5,s[1]) << '\n';
  cout << "sdev = " << fmt('f',10,5,s[2]) << '\n';

  s = simple(mv.geometric_risk_free_rate);

  cout << '\n';
  cout << "geometric_risk_free_rate" << '\n';
  cout << "mean = " << fmt('f',10,5,s[1]) << '\n';
  cout << "sdev = " << fmt('f',10,5,s[2]) << '\n';

  cout << '\n';

  cout << (rv ? "success" : "failure") << '\n';
*/
  return 0;
}

