#ifndef __FILE_MLEUSR_H_SEEN__
#define __FILE_MLEUSR_H_SEEN__
/* ----------------------------------------------------------------------------

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 "libmle.h"
#include "mle_base.h"

namespace mle {
    
  class habit_usrmod;

  typedef habit_usrmod usrmod_type;

  const INTEGER n_parms = 7;       //Number of parameters
  const INTEGER n_funcs = 5;       //g,sig,sig_w,rho,phi,delta,gamma
  const INTEGER n_data_rows = 8;   //year, gdp, inc, con, stk, sdiv, bnd, bdiv
  const INTEGER start_index = 9;
  const INTEGER start_year = 1950;
  const INTEGER HAC_lags = 1;

  struct habit_parameters {
    INTEGER p;                 // number of model parameters
    REAL g,sig;                // mean and sd. of log consumption growth
    REAL sig_w,rho;            // sd. of log dividend growth and corr. with v
    REAL phi;                  // surplus ratio autoregressive parameter
    REAL delta,gamma;          // time preference and risk aversion parmameters
    scl::realmat R;            // Cholesky factor of variance matix
    habit_parameters() 
    : p(n_parms)
    {                          // These are the values from Campbell & Cochrane
      scl::realmat theta(p,1); // They define the location of the prior
      theta[1] = 0.0189;       // g =     0.0189
      theta[2] = 0.015;        // sig =   0.015
      theta[3] = 0.112;        // sig_w = 0.112  
      theta[4] = 0.2;          // rho =   0.2
      theta[5] = 0.87;         // phi =   0.87
      theta[6] = 0.89;         // delta = 0.89
      theta[7] = 2.00;         // gamma = 2
      set_theta(theta);
    }
    void set_theta(const scl::realmat& theta)
    {
      if (theta.nrow() != p || theta.ncol() != 1) 
        scl::error("Error, habit_parameters, set_theta, bad dimension");
      R.resize(2,2);
      g     = theta[1];
      sig   = theta[2];
      sig_w = theta[3];
      rho   = theta[4];
      phi   = theta[5];
      delta = theta[6];
      gamma = theta[7];
    }
    void get_theta(scl::realmat& theta) const
    {
      theta.resize(p,1);
      theta[1] = g;
      theta[2] = sig;
      theta[3] = sig_w;
      theta[4] = rho;
      theta[5] = phi;
      theta[6] = delta;
      theta[7] = gamma;
    }
  };

  struct habit_variables {

    bool fixed_part_set;

    INTEGER n;   // sample size

    scl::realmat year;

    scl::realmat log_consumption_growth;     // Use to estimate g, sig, rho
    scl::realmat log_dividend_growth;        // Use to estimage g, sig_w, rho
                                             // Use lag as instrument
    scl::realmat log_surplus_ratio;
    scl::realmat sensitivity_function;

    scl::realmat consumption_growth;
    scl::realmat surplus_ratio;

    scl::realmat marginal_rate_of_substitution;

    scl::realmat gross_stock_return;
    scl::realmat gross_bond_return;

    scl::realmat geometric_stock_return;     // Use lag as instrument
    scl::realmat geometric_bond_return;      // Use lag as instrument

    scl::realmat log_gdp_growth;             // Too correlated with income
    scl::realmat log_income_growth;          // Use lag as instrument

    scl::realmat price_dividend_ratio; 
    scl::realmat log_price_dividend_ratio;   // Use as instrument

    habit_variables() 
    : fixed_part_set(false),
      n(0) 
    { }
    void resize_habit_variables(INTEGER sample_size);
    bool set_fixed_part (
      const scl::realmat& yyyy, const scl::realmat& gdp, 
      const scl::realmat& inc, const scl::realmat& con,
      const scl::realmat& stk, const scl::realmat& sdiv, 
      const scl::realmat& bnd, const scl::realmat& bdiv);
    bool set_variable_part(habit_parameters hp);
    std::vector<std::string> get_habit_variables(scl::realmat& hv);
  };

  class habit_moment_function : public scl::moment_function_base {
  private:
    const scl::realmat* dat;
    INTEGER sample_size;
    INTEGER L;
    scl::realmat theta;
    habit_parameters hp;
    habit_variables hv;
  public:
    habit_moment_function() : dat(0), sample_size(0), L(0), theta() { }
    bool set_data(const scl::realmat* datptr); 
    bool set_sample_size(INTEGER n); 
    bool set_L(INTEGER lag); 
    bool set_theta(const scl::realmat& parm); 
    INTEGER get_T0(); 
    INTEGER get_d();
    scl::realmat operator() (INTEGER t); 
    habit_parameters get_habit_parameters() {return hp;}
    habit_variables get_habit_variables() {return hv;} 
  };

  class habit_usrmod : public libmle::usrmod_base {
  private:
    scl::realmat data;
    INTEGER mflags;
    INTEGER haclags;
    scl::realmat prior_scale;
    bool use_jacobian;
    scl::realmat theta;
    habit_moment_function hmf;
    scl::gmm hgmm;
  public:
    habit_usrmod
      (const scl::realmat& dat, INTEGER len_mod_parm, INTEGER len_mod_func,
       const std::vector<std::string>& mod_pfvec,
       const std::vector<std::string>& mod_alvec, 
       std::ostream& detail);
    INTEGER len_rho() {return n_parms;}
    INTEGER len_stats() {return n_funcs;}
    bool get_stats(scl::realmat& stats);
    void get_rho(scl::realmat& parm) {parm = theta;}
    void set_rho(const scl::realmat& parm) {theta=parm; hmf.set_theta(parm);}
    bool support(const scl::realmat& rho);
    scl::den_val prior(const scl::realmat& rho, const scl::realmat& stats);
    scl::den_val likelihood();
    void write_usrvar(const char* filename); 
  };

}
#endif
