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

using namespace std;
using namespace scl;

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

  const INTEGER p_ge = 4;
  const INTEGER p_pe = 2;
  const INTEGER K_pe = 2;

  INT_32BIT seed = 740726;

  realmat prior_mean(p_ge,1);
  realmat prior_sdev(p_ge,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;

  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(p_ge,1);
  realmat theta(p_pe,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;

  const REAL logoneontwopi = log(1.0/6.283195307179587);

  realmat data(K_pe,n_obs);
  realmat data_orig(K_pe,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];
  }

  tree_moment_function tree_mf;

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

  success = tree_mf.set_sample_size(n_obs);
  if (!success) rv = false;

  success = tree_mf.set_L(mf_lags);
  if (!success) rv = false;

  success = tree_mf.set_theta(theta);
  if (!success) rv = false;

  INTEGER T0 = tree_mf.get_T0();

  INTEGER d = tree_mf.get_d();

  realmat mdata(n_obs,d,0.0);
  for (INTEGER t=T0+1; t<=n_obs; ++t) {
    realmat mt = tree_mf(t);
    for (INTEGER j=1; j<=d; ++j) mdata(t,j) = mt[j];
  }
  
  intvec midx=seq(1,d);
  realmat b0_var, B_var, S_var, C_var, V_var, T_var;
  varcoef(mdata,midx,1,b0_var,B_var,S_var,C_var);
  V_var = kronprd(C_var,S_var);
  T_var = cbind(b0_var,B_var);
  for (INTEGER i=1; i<=T_var.size(); ++i) T_var[i] /= sqrt(V_var(i,i));

  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);
  tree_gmm.set_regularize_W(true);

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

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

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

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

  d = tree_gmm.get_d();

  REAL log_likelihood = -0.5*REAL(n_obs)*obj + 0.5*REAL(d)*logoneontwopi;

  realmat R;
  realmat Rinv;

  //if (cholesky(W,R) != d) rv = false;  Do not use this form.
  //realmat z = sqrt(n_obs)*R*m;         Appears inaccurate for some reason.

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

  tree_mf.set_theta(theta);
  
  realmat U_svd,S_svd,V_svd;
  svd(S,U_svd,S_svd,V_svd);

  realmat z_svd = sqrt(n_obs)*T(U_svd)*m;
  for (INTEGER i=1; i<=d; ++i) z_svd[i] /= sqrt(S_svd[i]);

  realmat pricing_errors_alt(n_obs,1,0.0);
  realmat m_alt(d,1,0.0);
  for (INTEGER t=T0; t<=n_obs; ++t) {
     realmat m_t = tree_mf(t);
     m_alt += m_t;
     pricing_errors_alt[t] = m_t[1];
   }
   m_alt = m_alt/n_obs;

  realmat z_gmm;
  REAL log_likelihood_gmm = tree_gmm.likelihood(theta, z_gmm).log_den;

  cout << '\n';
  cout << "First tree_gmm call" << '\n';
  cout << "m " << m << '\n';
  cout << "m_alt " << m_alt << '\n';
  cout << "z " << z << '\n';
  cout << "z_gmm " << z_gmm << '\n';
  cout << "z_svd " << z_svd << '\n';
  cout << "S " << S << '\n';
  cout << "W " << W << '\n';
  cout << "Singular values    " << S_svd << '\n';
  cout << "T0                 " << T0 <<'\n';  
  cout << "Log_10 cond num    " << log10(S_svd[1]/S_svd[2]) << '\n';
  cout << "logdetW            " << logdetW << '\n';
  cout << "rankW              " << rankW << '\n';
  cout << "obj                " << obj << '\n';
  cout << "n*obj              " << n_obs*obj << '\n';
  cout << "z'z                " << (T(z)*z)[1] << '\n';
  cout << "log_likelihood     " << log_likelihood << '\n';
  cout << "log_likelihood_gmm " << log_likelihood_gmm << '\n';
  cout << '\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';

  cout << "b0_var, B_var, S_var, C_var, T_var for moments" 
        << b0_var << B_var << S_var << C_var << T_var;
  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);

  realmat prior_range(p_ge,2);
  for (INTEGER i=1; i<=p_ge; ++i) {
    prior_range(i,1) =  REAL_MAX;
    prior_range(i,2) = -REAL_MAX;
  }

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

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

    for (INTEGER i=1; i<=p_ge; ++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];

    for (INTEGER i=1; i<=p_ge; ++i) {
      prior_range(i,1) = rho[i] < prior_range(i,1) ? rho[i] : prior_range(i,1);
      prior_range(i,2) = rho[i] > prior_range(i,2) ? rho[i] : prior_range(i,2);
    }

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

    if (data != data) cerr << data;

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

  string filename = "Z_" + fmt('d',4,n_obs)('0') + ".txt";

  writetable(filename.c_str(),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 << "Prior range " << prior_range;

  cout << '\n';

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

  return 0;
}

