#include "tree_sim.h"
#include "tree_policy.h"

using namespace std;
using namespace scl;

bool tree_sim
  (const realmat& rho, INTEGER n_sim, INT_32BIT& seed, tree_variables& mv)
{

  /* 
  Values used for pc_calib.txt
  alpha    0.95000
  sigma    0.03000
  beta     0.95000
  gamma   12.60011 = 1.2600114143005971e+01
  */
  
  REAL alpha = rho[1];
  REAL sigma = rho[2];
  REAL beta  = rho[3];
  REAL gamma = rho[4];

  spline_interpolator P;
  spline_interpolator Q;
  spline_interpolator p;
  spline_interpolator q;

  mv.resize(n_sim);

  bool converge = tree_policy(alpha, sigma, beta, gamma, P, Q, p, q);

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

  for (INTEGER i=1; i<=100; ++i) {
    REAL y = alpha*y_lag + sigma*unsk(seed);
    y_lag = y;
  }

  for (INTEGER i=1; i<=mv.N; ++i) {
    REAL y = alpha*y_lag + sigma*unsk(seed);
    mv.log_consumption[i] = y;
    mv.log_stock_price[i] = p(y);
    REAL stock_payoff = exp(y) + exp(p(y));
    mv.log_marginal_rate_of_substitution[i] = log(beta) - gamma*(y - y_lag);
    mv.log_consumption_growth[i] = y - y_lag;
    mv.geometric_stock_return[i] = log(stock_payoff) - p(y_lag);
    mv.geometric_risk_free_rate[i] 
      = -log(beta) - (1.0-alpha)*gamma*y - 0.5*pow(gamma*sigma,2);
    mv.pricing_errors[i]
      = 1.0 - exp(mv.log_marginal_rate_of_substitution[i] +
        mv.geometric_stock_return[i]);
    y_lag = y;	
  }

  /*
  realmat pc_calib = cbind(mv.log_stock_price, mv.log_consumption);
  writetable("pc_calib.txt",pc_calib("1:1000",""),20,16);
  */

  return converge;
}

std::vector<std::string> 
tree_variables::get_tree_variables(scl::realmat& vars)
{
  std::vector<std::string> names(7);
  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);
  names[6] = "pricing_errors";
  vars=cbind(vars,pricing_errors);
  return names;
}
