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

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

#include <map>

using namespace std;
using namespace scl;

const bool debug = true;

// Raw sim is DCFdatayp.dat, M=5, n=86

// Col 1: Year: 1930 -- 2015
// Col 2: log of extracted mrs
// Col 3: growth of Real per capita GDP
// Col 4: GDP/Corporate profits
// Col 5: growth of Corporate profits

// Data for use is sim_orig, n=86, M=3

// Row 1: log of extracted mrs
// Row 2: log GDP/Corporate profits
// Row 3: log growth of Corporate profits

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 = 2000;
  const INTEGER L = 1;
  const INTEGER M = 3;
  const INTEGER years = 30;
  const INTEGER p_rho = M + M*M + (M*M+M)/2 + 2;

  const INTEGER n_raw_items = 5;  //year, mrs, gdp, gdp_cp, cp
  const INTEGER n_raw_obs   = 86; //Annual 1930 to 2015

  const INTEGER loc_dat_yr     = 1;  
  const INTEGER loc_dat_mrs    = 2;
  const INTEGER loc_dat_gdp    = 3;
  const INTEGER loc_dat_gdp_cp = 4;
  const INTEGER loc_dat_cp     = 5;

  const INTEGER loc_var_mrs    = 1;
  const INTEGER loc_var_gdp_cp = 2;
  const INTEGER loc_var_cp     = 3;

  const INTEGER mrs_pos = loc_var_mrs;
  const INTEGER cf_pos = loc_var_cp;

  realmat sim_raw;
  INTEGER ct = n_raw_items*n_raw_obs; 

  if (ct!=readtable("DCFdatayp.dat",sim_raw)) error("Error, readtable failed");

  realmat sim_orig(M,n_raw_obs);
  for (INTEGER j=1; j<=n_raw_obs; ++j) {
    sim_orig(loc_var_mrs,j)    = sim_raw(j,loc_dat_mrs);
    sim_orig(loc_var_gdp_cp,j) = sim_raw(j,loc_dat_gdp_cp);
    sim_orig(loc_var_cp,j)     = sim_raw(j,loc_dat_cp);
  }

  realmat y0(M,1);
  for (INTEGER i=1; i<=M; ++i) y0[i] = sim_orig(i,n_raw_obs);

  if (debug) {
    cout << '\n';
    cout << "First 6 observations: mrs,gdp_cp,cp " << sim_orig("",seq(1,6));
    cout << '\n';
    cout << "Last 6 observations " << sim_orig("",seq(n_raw_obs-5,n_raw_obs));
    cout << '\n';
    cout << "y0 " << y0 << '\n';
    cout << '\n';
    cout.flush();
  }
  
  realmat rho_orig, hess_orig;

  if(!vecread("dcf.rho_mode.dat",rho_orig)) error("Error, vecread failed");
  if(!vecread("dcf.V_hat_hess.dat",hess_orig)) error("Error, vecread failed");

  /*
  REAL hess_19_save = hess_orig(19,19);
  REAL hess_20_save = hess_orig(20,20);
  for (INTEGER i=1; i<=p_rho; ++i) {
     hess_orig(19,i)=hess_orig(i,19)=hess_orig(20,i)=hess_orig(i,20)=0.0;
  } 
  hess_orig(19,19) = hess_19_save/1000.0;
  hess_orig(20,20) = hess_20_save/1000.0;
  */

  realmat R_orig = hess_orig;
  if(factor(R_orig)!=0) error("Error, factor failed");

  if (debug) {
    cout << '\n';
    cout << "rho_orig = " << rho_orig << '\n';
    cout << "hess_orig = " << hess_orig << '\n';
    //cout << "R_orig = " << R_orig << '\n';
    cout << '\n';
    cout.flush();
  } 

  realmat rho = rho_orig;

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

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

  if (debug) {
    cout << '\n';
    cout << "b0 = " << b0 << '\n';
    cout << "B = " << B << '\n';
    cout << "R = " << R << '\n';
    cout << "S = " << S << '\n';
    cout << "theta = " << theta << '\n';
  }

  for (INTEGER i=M+M*M+1; i<=M+M*M+(M*(M+1))/2; ++i) {
    rho_orig[i] /= 10.0;
  }

  for (INTEGER i=M+M*M+1; i<=M+M*M+(M*(M+1))/2+2; ++i) {
    for (INTEGER j=1; j<=M+M*M+(M*(M+1))/2+2; ++j) {
      hess_orig(i,j) /= 10.0;
      hess_orig(j,i) /= 10.0;
    }
  }

  if (debug) {
    cout << '\n';
    cout << "rho_orig = " << rho_orig << '\n';
    cout << "hess_orig = " << hess_orig << '\n';
    //cout << "R_orig = " << R_orig << '\n';
    cout << '\n';
    cout.flush();
  } 

  rho = rho_orig;

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

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

  if (debug) {
    cout << '\n';
    cout << "b0 = " << b0 << '\n';
    cout << "B = " << B << '\n';
    cout << "R = " << R << '\n';
    cout << "S = " << S << '\n';
    cout << "theta = " << theta << '\n';
  }

  realmat ecf, pvcf, pv1, cumecf, dcf, yld;
  realmat ylag = y0;

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

  INT_32BIT seed = 893475;

  realmat sim(M,n_obs);
  realmat data(2,n_obs);

  realmat y(M,1), u(M,1);

  string filename = "lsrlcg.txt";
  ofstream fout;
  fout.open(filename.c_str());
  if (!fout) error("Error, cannot open " + filename);

  for (INTEGER t=1; t<=5000; ++t) {
    for (INTEGER i=1; i<=M; ++i) u[i] = unsk(seed);
    y = b0 + B*ylag + R*u;
    dcfyld(b0,B,S,y,mrs_pos,cf_pos,years,ecf,pvcf,pv1,cumecf,dcf,yld);
    REAL price = dcf[years];
    REAL lsr = log(price+y[cf_pos]) - log(price_lag);
    REAL lcg = (log(beta) - y[mrs_pos])/gamma;
    fout << fmt('f',22,16,lsr) << fmt('f',22,16,lcg) << '\n';
    if (t<=n_obs) {
      for (INTEGER i=1; i<=M; ++i) sim(i,t) = y[i];
      data(1,t) = lsr;
      data(2,t) = lcg;
    }
    ylag = y;
    price_lag = price;
  }

  fout.clear(); fout.close();

  if (debug) {
    cout << '\n';
    cout << "First 6 sim: mrs,gdp_cp,cp " << sim("",seq(1,6));
    cout << '\n';
    cout << "Last 6 sim: mrs,gdp_cp,cp " << sim("",seq(n_obs-5,n_obs));
    cout << '\n';
    cout << "First 6 data: lsr, lcg " << data("",seq(1,6));
    cout << '\n';
    cout << "Last 6 data: lsr, lcg " << data("",seq(n_obs-5,n_obs));
    cout << '\n';
    cout.flush();
  }

  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;

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

  if (debug) {
    cout << '\n';
    cout << "d " << d << '\n';
    cout << '\n';
    cout << "likelihood " 
         << "("<< likelihood.positive <<", "<< likelihood.log_den << ")"<<'\n';
    cout << '\n';
    cout << "z " << z << '\n';
    cout << '\n';
  }

  if (p_rho != rho.size()) error("Error, p_rho !=rho.size()");

  realmat e(p_rho,1);

  realmat Z_save(n_prior_reps,d);
  realmat rho_save(n_prior_reps,p_rho);

  realmat rho_left(p_rho,1,REAL_MAX);
  realmat rho_rite(p_rho,1,-REAL_MAX);

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

    bool bad_draw = false;
    do {
      for (INTEGER i=1; i<=p_rho; ++i) e[i] = unsk(seed);
      rho = rho_orig + R_orig*e;

      frac_rho(M, L, rho, b0, B, S, R, theta);
      beta = theta[1];
      gamma = theta[2];

      bad_draw = false;
      if ( (beta  <= 0.8) || (0.99 <= beta  ) ) bad_draw = true;
      if ( (gamma <= 0.0) || (100.0 <= gamma) ) bad_draw = true;

    } while (bad_draw);

    if (debug) if (rep%100 == 0) cout << rep << '\n';

    for (INTEGER i=1; i<=p_rho; ++i) {
      rho_save(rep,i) = rho[i];
      if (rho[i] < rho_left[i]) rho_left[i] = rho[i];
      if (rho[i] > rho_rite[i]) rho_rite[i] = rho[i];
    }

    ylag = y0;
    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) u[i] = unsk(seed);
      y = b0 + B*ylag + R*u;
      for (INTEGER i=1; i<=M; ++i) sim(i,t) = y[i];
      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;
    }

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

    likelihood = crra_gmm.likelihood(theta,z);

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

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

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

  writetable(filename.c_str(),Z_save,header,25,16);

  map<INTEGER,string> rho_names;
  rho_names[ 1] = "$b_0(1)$";
  rho_names[ 2] = "$b_0(2)$";
  rho_names[ 3] = "$b_0(3)$";
  rho_names[ 4] = "$B(1,1)$";
  rho_names[ 5] = "$B(2,1)$";
  rho_names[ 6] = "$B(3,1)$";
  rho_names[ 7] = "$B(1,2)$";
  rho_names[ 8] = "$B(2,2)$";
  rho_names[ 9] = "$B(3,2)$";
  rho_names[10] = "$B(1,3)$";
  rho_names[11] = "$B(2,3)$";
  rho_names[12] = "$B(3,3)$";
  rho_names[13] = "$R(1,1)$";
  rho_names[14] = "$R(1,2)$";
  rho_names[15] = "$R(2,2)$";
  rho_names[16] = "$R(1,3)$";
  rho_names[17] = "$R(2,3)$";
  rho_names[18] = "$R(3,3)$";
  rho_names[19] = "$\\beta$";
  rho_names[20] = "$\\gamma$";
  
  filename = "dcf_prior.tex";
  fout.open(filename.c_str());
  if (!fout) error("Error, cannot open " + filename);

  realmat prior_mean;
  realmat prior_var;

  varcov(rho_save,prior_mean,prior_var);

  if(!vecwrite("prior_mean",prior_mean))  error("Error, vecwrite failed");
  if(!vecwrite("prior_var",prior_var))    error("Error, vecwrite failed");

  for (INTEGER i=1; i<=p_rho; ++i) {
    fout << rho_names[i] 
         <<" &" 
         <<" & "<< fmt('f',6,4,rho_orig[i])
         <<" & "<< fmt('f',6,4,sqrt(hess_orig(i,i)))
         <<" &"; 
    if (i == 19) {
      fout <<" & " << "$0.8$" <<" & " << "$0.99$"; 
    }
    else if (i == 20) {
      fout <<" & " << "$0$" <<" & " << "$100$"; 
    }
    else {
      fout <<" & " << "$-\\infty$" <<" & " << "$\\infty$";
     }
    fout <<" &" 
         <<" & "<< fmt('f',6,4,rho_left[i])
         <<" & "<< fmt('f',6,4,rho_rite[i])
         <<" \\\\" << '\n';
  }

  fout.clear(); fout.close();

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

  return 0;
}

