#include "libscl.h"
#include "spline_interpolator.h"
#include "tree_policy.h"

using namespace std;
using namespace scl;

struct tree_usrmod_variables {
  INTEGER N;   // simulation length
  scl::realmat log_consumption;
  scl::realmat log_stock_price;
  scl::realmat log_marginal_rate_of_substitution;
  scl::realmat log_consumption_growth;
  scl::realmat geometric_stock_return;
  scl::realmat geometric_risk_free_rate;
  tree_usrmod_variables() { }
  tree_usrmod_variables(INTEGER num_sim)
  : N(num_sim),
    log_consumption(N,1,0.0),
    log_stock_price(N,1,0.0),
    log_marginal_rate_of_substitution(N,1,0.0),
    log_consumption_growth(N,1,0.0),
    geometric_stock_return(N,1,0.0),
    geometric_risk_free_rate(N,1,0.0)
  { }
  std::vector<std::string> get_tree_usrmod_variables(scl::realmat& mv);
};

vector<string> 
tree_usrmod_variables::get_tree_usrmod_variables(realmat& vars)
{
  vector<string> names(6);
  names[0] = "log_consumption";
  vars = log_consumption;
  names[1] = "log_stock_price";
  vars = cbind(vars,log_stock_price);
  names[2] = "log_marginal_rate_of_substitution";
  vars = cbind(vars,log_marginal_rate_of_substitution);
  names[3] = "log_consumption_growth";
  vars = cbind(vars,log_consumption_growth);
  names[4] = "geometric_stock_return";
  vars=cbind(vars,geometric_stock_return);
  names[5] = "geometric_risk_free_rate";
  vars=cbind(vars,geometric_risk_free_rate);
  return names;
}

int main(int argc, char** argp, char** envp)
{
  INT_32BIT seed = 740726;

  const INTEGER n_sim = 200000; 

  REAL g     = 0.02;
  REAL alpha = 0.95;
  REAL sigma = 0.02;
  REAL beta  = 0.98;
  REAL gamma = 2.0;
  REAL delta = beta*exp(-gamma*g);

  cout << '\n';
  cout << "Parameters" << '\n';
  cout << "g     " << fmt('f',10,5,g) << '\n';
  cout << "alpha " << fmt('f',10,5,alpha) << '\n';
  cout << "sigma " << fmt('f',10,5,sigma) << '\n';
  cout << "beta  " << fmt('f',10,5,beta) << '\n';
  cout << "gamma " << fmt('f',10,5,gamma) << '\n';
  cout << "delta " << fmt('f',10,5,delta) << '\n';

  spline_interpolator H;
  spline_interpolator R;
  spline_interpolator h;
  spline_interpolator r;

  tree_usrmod_variables mv(n_sim);

  bool converge = tree_policy(alpha, sigma, delta, gamma, H, R, h, r);

  if (!converge) warn("Warning, tree_policy did not converge");

  REAL sdev_y = sigma/sqrt(1.0 - alpha*alpha);
  REAL y_lag = sdev_y*unsk(seed);

  REAL c_lag = y_lag;

  for (INTEGER t=1; t<=mv.N; ++t) {
    REAL y = alpha*y_lag + sigma*unsk(seed);
    REAL c = g*REAL(t) + y;
    mv.log_consumption[t] = c;
    mv.log_stock_price[t] = g*REAL(t+1) + h(y);
    mv.log_marginal_rate_of_substitution[t] = -gamma*(c - c_lag);
    mv.log_consumption_growth[t] = c - c_lag;
    REAL detrend_log_stock_price_lag = h(y_lag);
    REAL detrend_log_stock_payoff = log(exp(h(y)) + exp(y));
    mv.geometric_stock_return[t]  
      = g + detrend_log_stock_payoff - detrend_log_stock_price_lag;
    mv.geometric_risk_free_rate[t] 
      = gamma*g - log(beta) - (1.0-alpha)*gamma*y - 0.5*pow(gamma*sigma,2);
    y_lag = y;	
    c_lag = c;
  }

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

  realmat pc_growth = cbind(mv.log_stock_price,mv.log_consumption);

  writetable("pc_growth.txt",pc_growth("1:1000",""),20,16);

  return 0;
}

