/* ----------------------------------------------------------------------------

Copyright (C) 2014.

A. Ronald Gallant
Post Office Box 659
Chapel Hill NC 27514-0659
USA

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

-----------------------------------------------------------------------------*/

#include "libscl.h"
#include "mleusr.h"

using namespace scl;
using namespace std;

namespace mle {

  habit_usrmod::habit_usrmod
    (const realmat& dat, INTEGER len_mod_parm, INTEGER len_mod_func,
     const vector<string>& mod_pfvec, const vector<string>& mod_alvec,
     ostream& detail)
  : data(dat), mflags(2), haclags(HAC_lags), prior_scale(n_parms,1,0.10), 
    use_jacobian(false)
  {
    detail << starbox("/habit_usrmod constructor now controls print//");
    detail << '\n';

    if (len_mod_parm != n_parms) error("Error, habit_usrmod, bad len_mod_parm");
    if (len_mod_func != n_funcs) error("Error, habit_usrmod, bad len_mod_func");

    if (mod_alvec.size() != 10) {
      error("Error, habit_usrmod, prior_scale or use_jacobian not in parmfile");
    }
    prior_scale[1] = atof(mod_alvec[1].substr(0,12).c_str());
    prior_scale[2] = atof(mod_alvec[2].substr(0,12).c_str());
    prior_scale[3] = atof(mod_alvec[3].substr(0,12).c_str());
    prior_scale[4] = atof(mod_alvec[4].substr(0,12).c_str());
    prior_scale[5] = atof(mod_alvec[5].substr(0,12).c_str());
    prior_scale[6] = atof(mod_alvec[6].substr(0,12).c_str());
    prior_scale[7] = atof(mod_alvec[7].substr(0,12).c_str());
    use_jacobian = atoi(mod_alvec[2].substr(0,12).c_str());

    detail << boolalpha << '\n';
    detail << "\t prior_scale  = " << prior_scale << '\n';
    detail << "\t use_jacobian = " << use_jacobian << '\n';
    detail << '\n';

    habit_parameters hp;
    hp.get_theta(theta);
    detail << "\t theta set by habit_usrmod constructor to" << theta << '\n';

    bool rv;

    rv = hmf.set_data(&data);
    if (!rv) error("Error, habit_usrmod, hmf.set_data failed ");

    INTEGER r = data.nrow();
    INTEGER c = data.ncol();

    rv = hmf.set_sample_size(r);
    if (!rv) error("Error, habit_usrmod, hmf.set_sample_size failed ");
    detail << "\t hmf.sample_size set by habit_usrmod constructor to " 
           << r << '\n';

    rv = hmf.set_L(mflags);
    if (!rv) error("Error, habit_usrmod, hmf.set_L failed ");
    detail << "\t hmf.L set by habit_usrmod constructor to " << mflags <<'\n';

    gmm tmp_gmm(&hmf,mflags,&data,c,haclags);

    detail << "\t hgmm constructed with these values:" << '\n';
    detail << "\t\t moment_function_lags = " << mflags << '\n';
    detail << "\t\t data first element data(1,1) = " << data(1,1) << '\n';
    detail << "\t\t data last element data(r,c) = " << data(r,c) << '\n';
    detail << "\t\t r = " << r << '\n';
    detail << "\t\t c = " << c << '\n';
    detail << "\t\t sample size = " << c << '\n';
    detail << "\t\t Lhac = " << haclags << '\n';

    hgmm = tmp_gmm;
    
    /*
    rv = hgmm.set_moment_function(&hmf);
    if (!rv) error("Error, habit_usrmod, hgmm.set_moment_function failed");
    detail << "\t hgmm.set_moment_function set by habit_usrmod constructor" 
           << '\n';

    rv = hgmm.set_moment_function_lags(mflags);
    if (!rv) error("Error, habit_usrmod, hgmm.set_moment_function_lags failed");
    detail << "\t hgmm.set_moment_function_lags set by habit_usrmod "
           << "constructor" << '\n';

    rv = hgmm.set_data(&data);
    if (!rv) error("Error, habit_usrmod, hgmm.set_data failed");
    detail << "\t hgmm.set_data set by habit_usrmod constructor" << '\n';

    rv = hgmm.set_sample_size(r);
    if (!rv) error("Error, habit_usrmod, hgmm.set_sample_size failed");
    detail << "\t hgmm.sample_size set by habit_usrmod constructor to " 
           << r << '\n';

    rv = hgmm.set_Lhac(haclags);
    if (!rv) error("Error, habit_usrmod, hgmm.set_Lhac failed");
    detail << "\t hgmm.Lhac set by habit_usrmod constructor to " 
           << haclags << '\n';
    */

    const bool correct_W_for_mean = true;
    hgmm.set_correct_W_for_mean(correct_W_for_mean);
    detail << "\t hgmm.set_correct_W_for_mean set by habit_usrmod "
           << "constructor to " << boolalpha << correct_W_for_mean << '\n';

    const bool regularize_W = true;
    const REAL ridge = 0.0;
    hgmm.set_regularize_W(regularize_W,ridge);
    detail << "\t hgmm.set_regularize_W set by habit_usrmod "
           << "constructor to (" << regularize_W << ", " << ridge <<')'<<'\n';

    const bool warning_messages = true;
    hgmm.set_warning_messages(warning_messages);
    detail << "\t hgmm.set_warning_messages set by habit_usrmod "
           << "constructor to " << warning_messages << '\n';

    detail << starbox("/mle now controls print//");
    detail.flush();
  }

  bool habit_usrmod::get_stats(realmat& stats)
  {
    realmat X;
    vector<string> names = hmf.get_habit_variables().get_habit_variables(X);

    INTEGER T0 = hmf.get_T0();

    INTEGER r = X.nrow();
    INTEGER c = X.ncol();

    string mrs_name = "marginal_rate_of_substitution";

    INTEGER mrs_idx = 0;
    for (INTEGER i=1; i<=c; ++i) {
      if (mrs_name == names[i-1]) mrs_idx = i;
    }
    
    if (mrs_idx == 0) error("Error, get_stats, can't find mrs");

    realmat mrs(r-T0+1,1);
    for (INTEGER i=T0; i<=r; ++i) mrs[i-T0+1] = X(i,mrs_idx);

    r = mrs.nrow();

    mrs.sort(1);

    stats.resize(n_funcs,1);

    stats[1] = mrs[1];
    if (n_funcs > 2) {
      for (INTEGER i=2; i<n_funcs; ++i) {
        REAL quantile = REAL(r)/REAL(n_funcs - 1);
	INTEGER pos = INTEGER(quantile);
	pos = pos < 1 ? 1 : pos;
	pos = pos > n_funcs ? n_funcs : pos;
        stats[i] = mrs[pos];
	pos += INTEGER(quantile);
      } 
    }
    stats[n_funcs] = mrs[r];

    return true;
  }

  den_val habit_usrmod::likelihood()
  {
    realmat m;
    realmat W;
    REAL logdetW;
    INTEGER rankW;
    realmat S;

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

    INTEGER d = hmf.get_d();
    const REAL logoneontwopi = log(1.0/6.283195307179587);

    INTEGER n = data.ncol() - hmf.get_T0() + 1;

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

    if (use_jacobian) log_likelihood += 0.5*logdetW;

     /*
     cerr << logdetW <<' '<< fmt('e',20,16,obj) 
     <<' '<< fmt('e',20,16,log_likelihood) << '\n';
     */

    return denval(true, log_likelihood);
  }

  bool habit_usrmod::support(const realmat& theta)
  {
    const REAL large = 0.0001*REAL_MAX;

    realmat theta_lo(n_parms,1);

    theta_lo[1] =  -0.5;       // g;
    theta_lo[2] =   0.0;       // sig;
    theta_lo[3] =   0.0;       // sig_w;
    theta_lo[4] =  -1.0;       // rho;
    theta_lo[5] =  -1.0;       // phi;
    theta_lo[6] =   0.7;       // delta;
    theta_lo[7] =   1.0;       // gamma;

    realmat theta_hi(n_parms,1);

    theta_hi[1] =   0.5;       // g;
    theta_hi[2] = large;       // sig;
    theta_hi[3] = large;       // sig_w;
    theta_hi[4] =   1.0;       // rho;
    theta_hi[5] =   1.0;       // phi;
    theta_hi[6] =  1.05;       // delta;
    theta_hi[7] =  20.0;       // gamma;

    return ( (theta_lo < theta) && (theta < theta_hi) );
  }

  den_val habit_usrmod::prior(const realmat& theta, const realmat& stats)
  {
    const REAL minus_log_root_two_pi = -9.1893853320467278e-01;

    habit_parameters tmp_hp;

    realmat mean(tmp_hp.p,1,0.0);
    mean[1] = tmp_hp.g;
    mean[2] = tmp_hp.sig;
    mean[3] = tmp_hp.sig_w;
    mean[4] = tmp_hp.rho;
    mean[5] = tmp_hp.phi;
    mean[6] = tmp_hp.delta;
    mean[7] = tmp_hp.gamma;

    tmp_hp.set_theta(theta);

    realmat rho(tmp_hp.p,1);
    rho[1] = tmp_hp.g;
    rho[2] = tmp_hp.sig;
    rho[3] = tmp_hp.sig_w;
    rho[4] = tmp_hp.rho;
    rho[5] = tmp_hp.phi;
    rho[6] = tmp_hp.delta;
    rho[7] = tmp_hp.gamma;

    realmat sdev(tmp_hp.p,1,100.0);
    for (INTEGER i=1; i<=tmp_hp.p; ++i) {
      sdev[i] = fabs(mean[i])*prior_scale[i]/1.96;
    }

    den_val sum(true,0.0);

    for (INTEGER i=1; i<=tmp_hp.p; ++i) {
      REAL z = (rho[i] - mean[i])/sdev[i];
      REAL e = minus_log_root_two_pi - log(sdev[i]) - 0.5*pow(z,2);
      sum += den_val(true,e);
    }
    return sum;
    //return den_val(true,0.0);
  }

  void habit_usrmod::write_usrvar(const char* filename)
  {
    string stem = filename;

    string::size_type suffix = stem.find(string("dat"));
    if (suffix != std::string::npos) {
      stem = stem.substr(0, suffix);
    }
    
    ofstream fout;
    fout.open(filename);
    if (!fout) error("Error, usrvar, could not open " + string(filename));

    ofstream floop;

    realmat X;
    vector<string> names = hmf.get_habit_variables().get_habit_variables(X);

    INTEGER T0 = hmf.get_T0();

    INTEGER r = X.nrow();
    //INTEGER c = X.ncol();

    vector<string>::const_iterator itr;
    INTEGER count = 0;
    for(itr=names.begin(); itr!=names.end(); ++itr) {
      ++count;
      string foutname = stem + *itr + ".dat";
      floop.open(foutname.c_str());
      if (!floop) error("Error, usrvar, could not open " + foutname);
      REAL sum = 0.0;
      for (INTEGER i=T0; i<=r; ++i) {
        sum += X(i,count);
	INTEGER year = INTEGER(X(i,1));
        floop << fmt('d',4,year)() 
	      << ' ' << fmt('f',17,8,X(i,count))
	      << ' ' << fmt('e',27,17,X(i,count))() 
	      << '\n';
      }
      floop.clear(); floop.close();
      REAL mean = sum/(r-T0+1);
      sum = 0.0;
      for (INTEGER i=T0; i<=r; ++i) {
        sum = pow(X(i,count)-mean,2);
      }
      REAL var = sum/REAL(r-T0);
      REAL sdev = sqrt(var);
      if (itr != names.begin()) {
        fout << fmt('l',30,*itr)() 
             << ' ' << fmt('f',17,8,mean)() 
	     << ' ' << fmt('f',17,8,sdev)() << '\n';
      }
    }
    fout.clear(); fout.close();
  }

  vector<string> habit_variables::get_habit_variables(realmat& hv)
  {
    vector<string> names(16);
    names[0] = "year";
    names[1] = "log_consumption_growth";
    names[2] = "log_dividend_growth";
    names[3] = "log_surplus_ratio";
    names[4] = "sensitivity_function";
    names[5] = "consumption_growth";
    names[6] = "surplus_ratio";
    names[7] = "marginal_rate_of_substitution";
    names[8] = "gross_stock_return";
    names[9] = "gross_bond_return";
    names[10] = "geometric_stock_return";
    names[11] = "geometric_bond_return";
    names[12] = "log_gdp_growth";
    names[13] = "log_income_growth";
    names[14] = "price_dividend_ratio";
    names[15] = "log_price_dividend_ratio";
    hv = year;
    hv = cbind(hv,log_consumption_growth);
    hv = cbind(hv,log_dividend_growth);
    hv = cbind(hv,log_surplus_ratio);
    hv = cbind(hv,sensitivity_function);
    hv = cbind(hv,consumption_growth);
    hv = cbind(hv,surplus_ratio);
    hv = cbind(hv,marginal_rate_of_substitution);
    hv = cbind(hv,gross_stock_return);
    hv = cbind(hv,gross_bond_return);
    hv = cbind(hv,geometric_stock_return);
    hv = cbind(hv,geometric_bond_return);
    hv = cbind(hv,log_gdp_growth);
    hv = cbind(hv,log_income_growth);
    hv = cbind(hv,price_dividend_ratio);
    hv = cbind(hv,log_price_dividend_ratio);
    return names;
  }

  void habit_variables::resize_habit_variables(INTEGER sample_size)
  {
    fixed_part_set = false;
    n = sample_size;
    year.resize(n,1,0.0);
    log_consumption_growth.resize(n,1,0.0);
    log_dividend_growth.resize(n,1,0.0);
    log_surplus_ratio.resize(n,1,0.0);
    sensitivity_function.resize(n,1,0.0);
    consumption_growth.resize(n,1,0.0);
    surplus_ratio.resize(n,1,0.0);
    marginal_rate_of_substitution.resize(n,1,0.0);
    gross_stock_return.resize(n,1,0.0);
    gross_bond_return.resize(n,1,0.0);
    geometric_stock_return.resize(n,1,0.0);
    geometric_bond_return.resize(n,1,0.0);
    log_gdp_growth.resize(n,1,0.0);
    log_income_growth.resize(n,1,0.0);
    price_dividend_ratio.resize(n,1,0.0);
    log_price_dividend_ratio.resize(n,1,0.0);
  }

  bool habit_variables::set_fixed_part(
    const realmat& yyyy, const realmat& gdp, const realmat& inc, 
    const realmat& con, const realmat& stk, const realmat& sdiv, 
    const realmat& bnd, const realmat& bdiv)
  {
    if (n == 0) {
      error("Error, habit_variables, n == 0");
      return false;
    }
    if (yyyy.size() != n || gdp.size() != n || bdiv.size() != n) {
      error("Error, habit_variables, spot check on sizes failed");
    }
    year[1] = yyyy[1];
    price_dividend_ratio[1] = stk[1]/sdiv[1];
    log_price_dividend_ratio[1] = log(price_dividend_ratio[1]);
    for (INTEGER i=2; i<=n; ++i) {
      year[i] = yyyy[i];
      log_consumption_growth[i] = log(con[i]) - log(con[i-1]);
      log_dividend_growth[i] = log(sdiv[i]) - log(sdiv[i-1]);
      consumption_growth[i] = con[i]/con[i-1];
      gross_stock_return[i] = (stk[i]+sdiv[i])/stk[i-1];
      gross_bond_return[i] = (bnd[i]+bdiv[i])/bnd[i-1];
      geometric_stock_return[i] = log(gross_stock_return[i]);
      geometric_bond_return[i] = log(gross_bond_return[i]);
      log_gdp_growth[i] = log(gdp[i]) - log(gdp[i-1]);
      log_income_growth[i] = log(inc[i]) - log(inc[i-1]);
      price_dividend_ratio[i] = stk[i]/sdiv[i];
      log_price_dividend_ratio[i] = log(price_dividend_ratio[i]);
    }
    fixed_part_set = true;
    return true;
  }

  bool habit_variables::set_variable_part(habit_parameters hp)
  {
    if (!fixed_part_set) {
      error("Error, habit_variables, fixed part not set");
      return false;
    }
    REAL del_c;

    REAL s;
    REAL S_bar = hp.sig*sqrt(hp.gamma/(1.0 - hp.phi));
    REAL s_bar = log(S_bar);
    REAL s_max = s_bar + (0.5)*(1.0 - pow(S_bar,2));
    REAL s_lag = s_bar;
    REAL del_s;

    REAL lam;
    REAL lam_lag = 0.0;

    REAL mrs;

    REAL v;

    for (INTEGER i=2; i <=n; i++) {

      del_c = log_consumption_growth[i];

      v = del_c - hp.g;

      s = (1.0 - hp.phi)*s_bar + hp.phi*s_lag + lam_lag*v;
      del_s = s - s_lag;
      s_lag = s;

      log_surplus_ratio[i] = del_s;
      surplus_ratio[i] = exp(del_s);

      lam = (s < s_max) ? sqrt(1.0 - 2.0*(s-s_bar))/S_bar - 1.0 : 0.0;
      lam_lag = lam;

      sensitivity_function[i] = lam;

      mrs = hp.delta*exp(-hp.gamma*(del_s + del_c));

      marginal_rate_of_substitution[i] = mrs;
    }

    return true;
  }

  bool habit_moment_function::set_data(const realmat* datptr)
  {
    dat=datptr; 

    INTEGER r = dat->nrow();
    INTEGER c = dat->ncol();

    if (r != n_data_rows) {
      error("Error, habit_moment_function, set_data, num rows wrong");
    }

    if (INTEGER((*dat)(1,start_index)) != start_year || start_index > c) {
      error("Error, habit_moment_function, set_data, years are wrong");
    }

    hv.resize_habit_variables(c);

    realmat year(c,1,0.0);
    realmat gdp(c,1,0.0);
    realmat inc(c,1,0.0);
    realmat con(c,1,0.0);
    realmat stk(c,1,0.0);
    realmat sdiv(c,1,0.0);
    realmat bnd(c,1,0.0);
    realmat bdiv(c,1,0.0);

    for (INTEGER j=1; j<=c; ++j) {
      year[j] = (*dat)(1,j);
      gdp[j]  = (*dat)(2,j);
      inc[j]  = (*dat)(3,j);
      con[j]  = (*dat)(4,j);
      stk[j]  = (*dat)(5,j);
      sdiv[j] = (*dat)(6,j);
      bnd[j]  = (*dat)(7,j);
      bdiv[j] = (*dat)(8,j);
    }

    bool rv = hv.set_fixed_part(year,gdp,inc,con,stk,sdiv,bnd,bdiv);
    if (!rv) {
      error("Error, habit_moment_function, set_data, set_fixed_part fail");
    }

    if (sample_size > c) {
      sample_size = c;
      warn("Warning, habit_moment_function, set_data, sample_size reset");
    }

    realmat theta_test;
    hp.get_theta(theta_test);

    if (theta != theta_test) {
      hp.set_theta(theta_test);
      theta = theta_test;
      warn("Warning, habit_moment_function, set_data, hp.theta reset");
    }
     
    rv = hv.set_variable_part(hp);
    if (!rv) {
      error("Error, habit_moment_function, set_data, set_variable_part fail");
    }
    
    return true;
  }

  bool habit_moment_function::set_sample_size(INTEGER n) 
  {
    if (n > hv.n) {
      error("Error, habit_moment_function, set_sample_size, n too large");
	return false;
    }
    sample_size = n; 
    return true;
  }

  bool habit_moment_function::set_L(INTEGER lag) 
  {
    if (lag > hv.n || lag < 2) {
	L = 2;
      warn("Warning, habit_moment_function, set_L, L set to 2");
      return false;
    }
    L = lag; 
    return true;
  }

  bool habit_moment_function::set_theta(const realmat& parm) 
  {
    theta = parm; 
    if (parm.size() != n_parms) {
      error("Error, habit_moment_function, set_theta, wrong size");
	return false;
    }
    hp.set_theta(theta);

    bool rv;
    rv = hv.set_variable_part(hp);
   
    if (!rv) {
      error("Error, habit_moment_function, set_data, set_variable_part fail");
    }

    return rv;
  }

  INTEGER habit_moment_function::get_T0() 
  {
    if (INTEGER(hv.year[start_index]) != start_year || start_index >= hv.n) {
      error("Error, habit_moment_function, set_data, hv.years are wrong");
    }
    return start_index;
  }

  INTEGER habit_moment_function::get_d() 
  {
    return 12;
  }

  realmat habit_moment_function::operator() (INTEGER t) 
  {
    realmat mt(12,1);
    REAL v = hv.log_consumption_growth[t] - hp.g;
    REAL w = hv.log_dividend_growth[t] - hp.g;
    mt[ 1] = v;
    mt[ 2] = pow(hp.sig,2) - pow(v,2);
    mt[ 3] = pow(hp.sig_w,2) - pow(w,2);
    mt[ 4] = hp.rho - v*w/(hp.sig*hp.sig_w);
    mt[ 5] = 1.0 - hv.marginal_rate_of_substitution[t]*hv.gross_stock_return[t];
    mt[ 6] = 1.0 - hv.marginal_rate_of_substitution[t]*hv.gross_bond_return[t];
    mt[ 7] = mt[5]*hv.geometric_stock_return[t-1];
    mt[ 8] = mt[5]*hv.geometric_bond_return[t-1];
    mt[ 9] = mt[5]*hv.log_income_growth[t-1];
    mt[10] = mt[6]*hv.geometric_stock_return[t-1];
    mt[11] = mt[6]*hv.geometric_bond_return[t-1];
    mt[12] = mt[6]*hv.log_income_growth[t-1];
    /*
    if (2006<hv.year[t] && hv.year[t]<2011) {
      cerr << hv.year[t] <<' '<< v 
           <<' '<< hv.sensitivity_function[t]
           <<' '<< hv.marginal_rate_of_substitution[t] << '\n';
    }
    */
    return mt;
  }
}
