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

Copyright (C) 2014, 2016, 2019

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

using namespace scl;
using namespace libmle;
using namespace mle;
using namespace std;

mle::crra_usrmod::crra_usrmod
  (const 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)
{ 
  detail << starbox("/crra_usrmod constructor now controls print//");
  detail << '\n';

  if (len_mod_parm != n_parms) error("Error, crra_usrmod, len_mod_parm wrong");
  if (len_mod_func != n_funcs) error("Error, crra_usrmod, len_mod_func wrong");

  detail << '\n';
  detail << "\t len_mod_parm = " << len_mod_parm << '\n'; 
  detail << "\t len_mod_func = " << len_mod_func << '\n'; 
  detail << '\n';
  detail.flush();

  
  parms.resize(n_parms,1);

  data = dat;

  if (mod_alvec.size() != 3) error("Error, crra_usrmod, bad mod_alvec");

  vector<string>::const_iterator al_ptr=mod_alvec.begin();

  ridge = atof((++al_ptr)->substr(0,12).c_str());

  detail << "\t ridge = "<< ridge <<'\n';

  bool rv;

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

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

  if (r != n_data_rows) error("Error, crra_usrmod, data has wrong nrow");
  if (c != n_data_cols) error("Error, crra_usrmod, data has wrong ncol");

  rv = crramf.set_sample_size(c);
  if (!rv) error("Error, crra_usrmod, crramf.set_sample_size failed ");

  if (crramf.get_d() != n_moments) error("Error, crra_usrmod, n_moments wrong");

  detail << "\t crramf.sample_size set by crra_usrmod constructor to "<<c<<'\n';

  INTEGER mflags = 1;

  rv = crramf.set_L(mflags);
  if (!rv) error("Error, crra, hmf.set_L failed ");

  gmm tmp_gmm(&crramf,mflags,&data,c,HAC_lags);

  detail << "\t crragmm 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 = " << HAC_lags << '\n';

  crragmm = tmp_gmm;

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

  const bool regularize_W = true;
  crragmm.set_regularize_W(regularize_W,ridge);
  detail << "\t crragmm.set_regularize_W set by crra_usrmod "
         << "constructor to (" << regularize_W << ", " << ridge <<')'<<'\n';

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

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


bool mle::crra_usrmod::support(const realmat& rho) 
{
  REAL gamma = parms[1];
  REAL delta = parms[2];

  if ( (delta <= 0.0) || (1 <= delta) ) return false;

  if (gamma <= 0.0) return false;

  return true;
}


den_val mle::crra_usrmod::prior(const realmat& rho, const realmat& stats) 
{
  //const REAL minus_log_root_two_pi = -log(sqrt(2.0*M_PI));
  const REAL minus_log_root_two_pi = -9.1893853320467267e-01;

  REAL logf = 0.0;

  REAL z = 0.0;
  REAL e = minus_log_root_two_pi;

  REAL gamma = parms[1];
  REAL delta = parms[2];

  const REAL gamma_mean = 10.0;
  const REAL gamma_sdev = 5.0;
  //const REAL gamma_sdev = 50.0;

  z = (gamma - gamma_mean)/gamma_sdev;
  e = minus_log_root_two_pi - log(gamma_sdev) - 0.5*pow(z,2);
  logf += e;

  const REAL delta_mean = 0.98039215686275;  // 2.0 per cent per annum
  const REAL delta_sdev = 0.02;
  
  z = (delta - delta_mean)/delta_sdev;
  e = minus_log_root_two_pi - log(delta_sdev) - 0.5*pow(z,2);
  logf += e;

  return den_val(true, logf);
}

den_val mle::crra_usrmod::likelihood() 
{
  realmat m;
  realmat W;
  REAL logdetW;
  INTEGER rankW;
  realmat S;
  
  REAL obj = crragmm(parms,m,W,logdetW,rankW,S);

  INTEGER d = crramf.get_d();

  bool positive = true;

  for (INTEGER i=1; i<=d; ++i) if (!IsFinite(m[i])) positive = false;

  if (!positive) {
    string msg = "Warning, get_likelihood, fail, moments not finite \n";
    msg += string("gamma, delta: ");
    msg += fmt('f',12,5,parms[1])();
    msg += fmt('f',12,5,parms[2])();
    msg += "\n";
    warn(msg);
  }

  if (!positive) return den_val(positive,-REAL_MAX);

  const REAL logoneontwopi = log(1.0/6.2831853071795865);
  //const REAL logoneontwopi = log(1.0/(2.0*M_PI));

  INTEGER n = data.ncol();

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

  return den_val(positive,log_likelihood);
}

void mle::crra_usrmod::write_usrvar(const char* filename)
{
  string stem;
  const char* fn = filename;
  while(*fn) stem.push_back(*fn++);
  string::size_type nchar = stem.size();
  stem = stem.substr(0,nchar-3);

  ofstream fout;
  string usrfile;

  usrfile = stem+"results.txt";
  fout.open(usrfile.c_str());
  if (!fout) error("Error, crra_usrmod, write_usrvar, cannot open " + usrfile);

  realmat m;
  realmat W;
  REAL logdetW;
  INTEGER rankW;
  realmat S;

  crragmm(parms,m,W,logdetW,rankW,S);

  realmat M = crragmm.get_M(parms);

  INTEGER n = data.ncol();
  INTEGER p = len_rho();

  realmat var = invpsd( REAL(n-crramf.get_T0()) * T(M)*W*M );

  realmat sdev(p,1);
  for (INTEGER i=1; i<=p; ++i) sdev[i] = sqrt(var(i,i));

  fout << "parms = " << parms << '\n';
  fout << "sdev = " << sdev << '\n';
  fout << "var = " << var << '\n';
}

