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

Copyright (C) 2013

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::elec_usrmod::elec_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)
: data(dat), rho(), lrho(11), lstats(2), a(), B(), R(),
  yhat(2,dat.ncol()), ehat(2,dat.ncol()), zhat(2,dat.ncol())
{
  if (lrho != len_mod_parm) {
    error("Error, usrmod, constructor, len_mod_parm is set wrong in parmfile");
  }

  if (lstats != len_mod_func) {
    error("Error, usrmod, constructor, len_mod_func is set wrong in parmfile");
  }
}

void mle::elec_usrmod::set_parms()
{
  a.resize(3,1);
  B.resize(3,3);
  R.resize(2,2);

  a[1] = rho[1];
  a[2] = rho[2];

  B(1,1) = rho[3];
  B(1,2) = rho[4];
  B(2,2) = rho[5];
  B(1,3) = rho[6];
  B(2,3) = rho[7];
  B(3,3) = rho[8];

  R(1,1) = rho[9];
  R(1,2) = rho[10];
  R(2,2) = rho[11];

  a[3] = -1.0;

  B(2,1) = B(1,2);
  B(3,1) = B(1,3);
  B(3,2) = B(2,3);

  R(2,1) = 0.0;
}

bool mle::elec_usrmod::get_obj(realmat& obj)
{
  INTEGER n = data.get_cols();

  yhat.resize(2,n,0.0);
  ehat.resize(2,n,0.0);
  zhat.resize(2,n,0.0);

  obj.resize(1,n,-REAL_MAX);

  if (!support(rho)) return false;

  realmat y = data("1:2","");
  realmat x = data("3:5","");

  realmat s = B*x;
  for (INTEGER t=1; t<=n; ++t) {
    s(1,t) += a[1];
    s(2,t) += a[2];
    s(3,t) += a[3];
  }

  for (INTEGER i=1; i<=s.size(); ++i) {
    if (s[i] >= 0.0) return false;
  }

  for (INTEGER t=1; t<=n; ++t) {
    REAL bot = log(-s(3,t));
    yhat(1,t) = log(-s(1,t)) - bot;
    yhat(2,t) = log(-s(2,t)) - bot;
    ehat(1,t) = y(1,t) - yhat(1,t);
    ehat(2,t) = y(2,t) - yhat(2,t);
  }

  realmat P = inv(R);

  zhat = P*ehat;

  REAL detR = R(1,1)*R(2,2);

  const REAL pi = 3.14159265358979312e+00;

  for (INTEGER t=1; t<=n; ++t) {
    obj[t] = pow(zhat(1,t),2) + pow(zhat(2,t),2);
    obj[t] *= (-0.5);
    obj[t] -= log(detR);
    obj[t] -= log(sqrt(2.0*pi));
  }

  return true;
}

den_val mle::elec_usrmod::likelihood() 
{
  INTEGER r = data.get_rows();
  INTEGER n = data.get_cols();

  if (r != 5) error("Error, elec_usrmod, likelihood, bad data");
  if (rho.get_rows() != 11) error("Error, elec_usrmod, likelihood, bad parm");

  realmat obj;

  if (!get_obj(obj)) return den_val(false,-REAL_MAX);;

  REAL q = 0.0;

  for (INTEGER t=1; t<=n; ++t) {
    q += obj[t];
  }

  return den_val(true,q);
}

bool mle::elec_usrmod::support(const realmat& parm) 
{
  if (parm[9] <= 0.0) return false;
  if (parm[11] <= 0.0) return false;
  return true; 
}

bool mle::elec_usrmod::get_stats(scl::realmat& stats) 
{ 
  stats.resize(2,1);
  INTEGER n = zhat.ncol();
  stats[1] = stats[2] = 0.0;
  for (INTEGER j=1; j<=n; ++j) {
    stats[1] += zhat(1,j);
    stats[2] += zhat(2,j);
  }
  stats[1] /=n; stats[2] /= n;
  return true;
}     

den_val mle::elec_usrmod::prior(const realmat& rho_in, const realmat& stats) 
{
  return den_val(true, 0.0);
}

bool mle::elec_usrmod::get_scores(scl::realmat& scores) 
{ 
  realmat rhosave = rho;

  INTEGER n = data.get_cols();

  scores.resize(lrho,n,0.0);

  REAL eps = std::pow(double(REAL_EPSILON),0.33333333);

  for (INTEGER i=1; i<=lrho; ++i) {
    REAL tmp = rho[i];

    REAL h = eps*std::fabs(tmp);
    if (h == 0.0) h = eps;

    REAL hi = tmp + h;
    rho[i] = hi;
    set_rho(rho);

    realmat f1;
    if (!get_obj(f1)) {set_rho(rhosave); return false;}

    REAL lo = tmp - h;
    rho[i] = lo;
    set_rho(rho);

    realmat f0;
    if (!get_obj(f0)) {set_rho(rhosave); return false;}

    REAL difference = hi - lo;

    for (INTEGER t=1; t<=n; ++t) {
      scores(i,t) = (f1[t] - f0[t])/difference;
    }

    rho[i] = tmp;
  }

  set_rho(rhosave);

  return true;
}
