#include "libscl.h"
#include "kronprd.h"

#include "var1EXX.h"
#include "lprior.h"
#include "crra_mf.h"

using namespace std;
using namespace scl;

const bool debug = true;

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 = 1000;

  // tau = (b0,B);
  // S = S
  // theta = (beta,gamma)

  const INTEGER L = 1;
  const INTEGER M = 3;
  const INTEGER p_tau = M+M*M;
  const INTEGER p_R = (M*M-M)/2;
  const INTEGER p_S = M*M;
  const INTEGER p_theta = 2;
  const INTEGER p_rho = p_tau + p_R + p_theta;
  const INTEGER mrs_pos = 1;
  const INTEGER cf_pos = 3;
  const INTEGER years = 30;

  const REAL tau_scale = 1.0;      // scale factor for tau_prior
  const INTEGER S_df = 10;         // degrees freedom of Wishart
  const REAL theta_scale = 1.0;    // scale factor for theta_prior

  // prior for (b0,B)

  realmat rho;

  if(vecread("dcf.rho_mode.dat", rho) == 0) error("Error, vecread failed");

  realmat b0, B, S, R, theta;
  frac_rho(M, L, rho, b0, B, S, R, theta);

  realmat rho_orig = rho;
  realmat b0_orig = b0;
  realmat B_orig = B;
  realmat S_orig = S;
  realmat R_orig = R;
  realmat theta_orig = theta;

  REAL beta = theta[1];
  REAL gamma = theta[2];

  REAL beta_orig = beta;
  REAL gamma_orig = gamma;

  realmat tau_mean(p_tau,1);
  for (INTEGER i=1; i<=p_tau; ++i) tau_mean[i] = rho[i]; 

  realmat EXX;
  var1EXX(b0, B, S, EXX);

  INTEGER rank;
  realmat C = invpsd(EXX,rank);
  if (rank != M+1) error("Error, EXX is not full rank");

  realmat tau_V = kronprd(C,S);
  realmat tau_R = tau_V;
  if (factor(tau_R)!=0) error("Error, factor failed");

  realmat E;
  INTEGER ier;
  REAL maxlam = eigen(B,E,ier);

  // prior for Sinv

  realmat S_parm = S;
  realmat S_parm_R = R;

  // prior for theta 

  realmat theta_mean(2,1);
  realmat theta_R(2,2);

  REAL beta_mean = theta_mean[1] = 0.95;
  REAL gamma_mean = theta_mean[2] = 12.5;

  REAL beta_sdev = 0.01;
  REAL gamma_sdev = 2.00;
  REAL beta_gamma_rho = 0.0;

  theta_R(1,1) = sqrt(pow(beta_sdev,2)*(1.0-pow(beta_gamma_rho,2)));
  theta_R(2,1) = 0.0;
  theta_R(1,2) = beta_sdev*beta_gamma_rho;
  theta_R(2,2) = gamma_sdev;

  realmat theta_V = theta_R*T(theta_R);

  cout << '\n';
  cout << "rho = rho_orig = " << rho_orig << '\n';
  cout << "b0 = b0_orig = " << b0_orig << '\n';
  cout << "B = B_orig = " << B_orig << '\n';
  cout << "R = R_orig = " << R_orig << '\n';
  cout << "S = S_orig = " << S_orig << '\n';
  cout << "theta = theta_orig = " << theta_orig << '\n';
  cout << "beta = beta_orig = " << beta_orig << '\n';
  cout << "gamma = gamma_orig = " << gamma_orig << '\n';
  cout << "EXX = " << EXX << '\n';
  cout << "C = " << C << '\n';
  cout << "tau_V = kronprd(C,S) = " << tau_V << '\n';
  cout << '\n';
  cout << "ier = " << ier << '\n';
  cout << "maxlam = " << maxlam << '\n';
  cout << '\n';
  cout << "tau_mean = " << tau_mean << '\n';
  cout << "tau_R = " << tau_R << '\n';
  cout << "tau_V = kronprd(C,S) = tau_R*T(tau_R) = " << tau_R*T(tau_R) << '\n';
  cout << "S_parm = " << S_parm << '\n';
  cout << "S_parm_R = " << S_parm_R << '\n';
  cout << "theta_mean = " << theta_mean << '\n';
  cout << "beta_mean = " << beta_mean << '\n';
  cout << "gamma_mean = " << gamma_mean << '\n';
  cout << "theta_R = " << theta_R << '\n';
  cout << "theta_V = " << theta_V << '\n';
  cout << '\n';

  INT_32BIT seed = 740726;
  realmat z;
  denval tau_dv, S_dv, theta_dv;

  // draw from tau prior
  
  realmat tau;

  z.resize(p_tau,1);
  bool stable = false;
  while (!stable) {
    for (INTEGER i=1; i<=p_tau; ++i) z[i] = unsk(seed);
    tau = tau_mean + tau_scale*tau_R*z;
    for (INTEGER i=1; i<=M*M; ++i) B[i] = tau[M+i];
    maxlam = eigen(B,E,ier);
    if (fabs(maxlam) < 1.0) stable = true;
  }
  tau_dv = mvn(tau,tau_mean,tau_V);
  for (INTEGER i=1; i<=M; ++i) b0[i] = tau[i];
  for (INTEGER i=1; i<=M*M; ++i) B[i] = tau[M+i];

  cout << boolalpha;
  cout << '\n';
  cout << "draw from tau prior = " << tau << '\n';
  cout << "b0 = " << b0 << '\n';
  cout << "B = " << B << '\n';
  cout << "density of tau draw = " 
       << '('<<tau_dv.positive<<','<<tau_dv.log_den<<')' << '\n';
  cout << "tau draw stable? " << stable << ", maxlam = " << maxlam << '\n';

  // draw from S prior

  z.resize(M,1);
  fill(S);
  for (INTEGER t=1; t<=S_df; ++t) {
    for (INTEGER i=1; i<=M; ++i) z[i] = unsk(seed);
    z = S_parm_R*z;
    S += z*T(z);
  }
  R = S;
  if (factor(R)!=0) error("Error, factor failed");
  S_dv = wishart(S_df, S, S_parm);

  cout << '\n';
  cout << "draw from S prior = " << S << '\n';
  cout << "R = " << R << '\n';
  cout << "density of S draw = " 
       << '('<<S_dv.positive<<','<<S_dv.log_den<<')' << '\n';

  // draw from theta prior
  
  z.resize(2,1);
  for (INTEGER i=1; i<=2; ++i) z[i] = unsk(seed);
  theta = theta_mean + theta_scale*theta_R*z;
  theta_dv = mvn(theta,theta_mean,theta_V);
  beta = theta[1];
  gamma = theta[2];
  cout << '\n';
  cout << "draw from theta prior = " << theta << '\n';
  cout << "beta = " << beta << '\n';
  cout << "gamma = " << gamma << '\n';
  cout << "density of theta draw = " 
       << '('<<theta_dv.positive<<','<<theta_dv.log_den<<')' << '\n';

  // stationary distribution

  realmat I(M,M,0.0);
  for (INTEGER i=1; i<=M; ++i) I(i,i) = 1.0;
  //realmat I = diag(realmat(M,1,1.0));

  realmat IB = I - B;
  realmat IBinv = inv(IB);

  realmat lim_mean = IBinv*b0;
  realmat vecS(M*M,1);
  for (INTEGER i=1; i<=M*M; ++i) vecS[i] = S[i];
  realmat vecV = inv(kronprd(I,I) - kronprd(B,B))*vecS;
  realmat lim_V(2,2);
  for (INTEGER i=1; i<=M*M; ++i) lim_V[i] = vecV[i];
  realmat lim_R = lim_V;
  if (factor(lim_R)!=0) error("Error, factor failed");

  cout << '\n';
  cout << "lim_mean = " << lim_mean << '\n';
  cout << "lim_V = " << lim_V << '\n';
  cout << "lim_R = " << lim_R << '\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';

  // simulate data

  realmat y, ecf, pvcf, pv1, cumecf, dcf, yld;

  realmat data(M,n_obs);

  z.resize(M,1);
  for (INTEGER i=1; i<=M; ++i) z[i] = unsk(seed);

  realmat ylag = lim_mean + lim_R*z;

  if (debug) {
    cout << '\n';
    cout << "b0, B, S, and R reset to original values" << '\n';
    cout << '\n';
    b0 = b0_orig;
    B  = B_orig;
    S  = S_orig;
    R  = R_orig;
    beta = beta_orig;
    gamma = gamma_orig;
    ylag = b0 + R*z;
  }

  for (INTEGER t=1; t<=100; ++t) {
    for (INTEGER i=1; i<=M; ++i) z[i] = unsk(seed);
    realmat y = b0 + B*ylag + R*z;
    ylag = y;
  }

  dcfyld(b0,B,S,ylag,mrs_pos,cf_pos,years,ecf,pvcf,pv1,cumecf,dcf,yld);
  REAL price_lag = dcf[years];

  for (INTEGER t=1; t<=n_obs; ++t) {
    for (INTEGER i=1; i<=M; ++i) z[i] = unsk(seed);
    realmat y = b0 + B*ylag + R*z;
    dcfyld(b0,B,S,y,mrs_pos,cf_pos,years,ecf,pvcf,pv1,cumecf,dcf,yld);
    REAL price = dcf[years];
    data(1,t) = log(price+y[cf_pos]) - log(price_lag);
    data(2,t) = (log(beta) - y[mrs_pos])/gamma;
    ylag = y;
    price_lag = price;
  }

  cout << simple(T(data));

  bool success;
  bool rv = true;

  mle::crra_moment_function crra_mf;

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

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

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

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

  INTEGER T0 = crra_mf.get_T0();

  INTEGER d = crra_mf.get_d();

  gmm crra_gmm(&crra_mf, mf_lags, &data, n_obs, HAC_lags);

  crra_gmm.set_correct_W_for_mean(true);
  crra_gmm.set_warning_messages(true);
  crra_gmm.set_regularize_W(true);

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

  denval likelihood = crra_gmm.likelihood(theta,z);

/*
  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 = crra_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 = crra_gmm.likelihood(theta, z_gmm).log_den;

  cout << '\n';
  cout << "First crra_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 = crra_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 = crra_gmm.set_data(&data);
    if (!success) rv = false;

    obj = crra_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;
}

