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

Copyright (C) 2009, 2011, 2012, 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.

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

#define USE_VARIABLE_SEED
//#undef USE_VARIABLE_SEED

#include <cerrno>
#include "libmle.h"
#include "mle.h"
#include "svparms.h"

using namespace scl;
using namespace std;

using namespace libmle;
using namespace mle;

namespace mle {

  liml_usrmod::liml_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)
  : data(dat),lrho(4),lstats(1),L(0),Lhac(0),N(2500),variable_seed(740726),
    saved_particle(),emcount(0),emlimit(50),gibbs_count(0)
  {
    detail << starbox("/Begin liml_usrmod constructor output//");

    if (len_mod_parm != lrho) error("Error, liml_usrmod, bad len_mod_parm");
    if (len_mod_func != lstats) error("Error, liml_usrmod, bad len_mod_parm");

    rho.resize(lrho,1);
    rho[1] = T1;     // rho
    rho[2] = T2;     // phi
    rho[3] = T3;     // sigma
    rho[4] = T4;     // beta, not estimated, must held fixed in parmfile

    rho_old = rho;

    gibbs_draws.reserve(5000);

    realmat svparticle;
    if (vecread("svparticle.dat",svparticle) == 0) {
      error("Error, liml_usrmod, cannot read svparticle.dat");
    }
    INTEGER n = data.ncol();
    INTEGER d = svparticle.nrow();
    saved_particle.resize(d,n+1);
    for (INTEGER t=1; t<=n+1; ++t) {
      for (INTEGER k=1; k<=d; ++k) {
        saved_particle(k,t) = svparticle(k,t);
      }
    }

    try  {
      mf = new svmnt();
      objf = new gmm(mf, L, &data, data.ncol(), Lhac);
      objf->set_regularize_W(true, 1.0e-3);
      m = new svmod(mf, objf);
    }
    catch (exception& e) {
      error("Error, liml_usrmod" + string(e.what()) );
    }

    detail << '\n';
    detail << "\t mf, objf, and m successfully instantiated" << '\n';

    if (d != 1) {
      error("Error, liml_usrmod, svparticle has wrong dimension");
    }


    if (mod_alvec.size() != 7) {
      error("Error, liml_usrmod, L, Lhac, N, slen not in parmfile");
    }
    L = atoi(mod_alvec[1].substr(0,12).c_str());
    Lhac = atoi(mod_alvec[2].substr(0,12).c_str());
    N = atoi(mod_alvec[3].substr(0,12).c_str());
    slen = atoi(mod_alvec[4].substr(0,12).c_str());
    emlimit = atoi(mod_alvec[5].substr(0,12).c_str());
    
    detail << '\n';
    detail << "\t    lrho = " << lrho << '\n';
    detail << "\t  lstats = " << lstats << '\n';
    detail << "\t       L = " << L << '\n';
    detail << "\t    Lhac = " << Lhac << '\n';
    detail << "\t       N = " << N << '\n';
    detail << "\t    slen = " << slen << '\n';
    detail << "\t emlimit = " << emlimit << '\n';
    detail << '\n';
    detail.flush();

    try  {
      draws = new vector<realmat>(N);
      smooth = new vector<realmat>(N);
      filter = new vector<realmat>(N);
    }
    catch (exception& e) {
      error("Error, liml_usrmod" + string(e.what()) );
    }

    detail << "\t draws, smooth, filter successfully instantiated with size "
           << (*draws).size()  << '\n';

    warn("Warning, USE_VARIABLE_SEED is defined");

    detail << starbox("/End liml_usrmod constructor output//");
    detail << '\n';
    detail.flush();
  }
  
  bool liml_usrmod::get_stats(realmat& stats)
  {
    INT_32BIT seed = fixed_seed;
    realmat sim = (*m).draw_sample(slen, seed).y;
    stats = realmat(1,1,0.0);
    INTEGER n = sim.size();
    for (INTEGER t=1; t<=n; ++t) stats[1] += sim[t];
    stats[1] /= REAL(n);
    return true;
  }
  
  bool liml_usrmod::support(const realmat& rho)
  {
    if (fabs(rho[1]) >= 0.999) return false;
    if (fabs(rho[2]) >= 0.999) return false;
    if (rho[3] <= 0.0) return false;
    return true;
  }
  
  denval liml_usrmod::prior(const realmat& theta, const realmat& stats)
  {
    return denval(true, 0.0);
  }

  denval liml_usrmod::likelihood()
  {

    #if defined USE_VARIABLE_SEED
      INT_32BIT seed = variable_seed;
    #else
      INT_32BIT seed = fixed_seed;
    #endif

    mf->set_theta(rho);
    m->set_parms(rho);

    INTEGER n = data.ncol();
    INTEGER T0 = mf->get_T0();

    realmat y(1,n+1);
    y[1] = data[1];
    for (INTEGER t=1; t<=n; ++t) y[1+t] = data[t];

    bool dset = objf->set_data(&y);          // sets in mf also
    bool nset = objf->set_sample_size(n+1);  // sets in mf also
    if (!(dset && nset)) error("Error, liml_usrmod, likelihood, config fail");

    REAL weights[N];
    REAL sum;

    realmat* partptr = &saved_particle;
    denval rv = m->likelihood(n+1,partptr);

    if (emcount < emlimit) {
      ++emcount;
      return rv;
    }

    emcount = 0;

    mf->set_theta(rho_old);
    m->set_parms(rho_old);

    gibbs_draws.push_back(rho_old);

    realmat x = saved_particle;

    // Initialization

    (*filter)[0] = (*draws)[0] = (*smooth)[0] = x;  

    for (INTEGER i=1; i<N; ++i) {
      x[1] = m->draw_x0(seed);
      (*filter)[i] = (*draws)[i] = (*smooth)[i] = x;
    }
  
    for (INTEGER t=2; t<T0; ++t) {
      for (INTEGER i=1; i<N; ++i) {
        (*filter)[i][t] = (*draws)[i][t] = (*smooth)[i][t] 
          = m->draw_xt((*draws)[i][t-1],seed);
      }
    }  
  
    for (INTEGER t=T0; t<=n+1; ++t) {
  
      vector<realmat>* save = draws;
      draws = smooth;
      smooth = save;
  
      // Importance sampling step
  
      partptr = &saved_particle;
      sum = weights[0] = m->prob_yt(t,partptr);
      for (INTEGER i=1; i<N; ++i) {
        realmat* partptr = &( (*draws)[i] );
        (*partptr)[t] = m->draw_xt((*partptr)[t-1],seed);
        sum += weights[i] = m->prob_yt(t,partptr);
      }

      weights[0] /= sum;
      for (INTEGER i=1; i<N; ++i) {
        weights[i] /= sum;
        weights[i] += weights[i-1];
      }
      weights[N-1] = 1.0;
  
      // Selection step
  
      for (INTEGER i=1; i<N; ++i) {
        REAL u = ran(seed);
        INTEGER j = 0;
        while(weights[j] <= u) ++j;
        (*smooth)[i] = (*draws)[j];
        (*filter)[i][t] = (*draws)[j][t];
      }

      for (INTEGER i=0; i<N; ++i) weights[i] = 1.0/REAL(N);
    }

    saved_particle = (*smooth)[N-1];

    #if defined USE_VARIABLE_SEED
      variable_seed = seed;
    #endif

    /*
    string filename = "svsim.usrvar.dat";
    write_usrvar(filename.c_str());
    */

    return rv;
  }

  void liml_usrmod::write_usrvar(const char* filename)
  {
    string stem = filename;
    size_t suffix = stem.find(string("dat"));
    if (suffix != std::string::npos) {
      stem = stem.substr(0, suffix);
    }

    ofstream fout;
    string foutname = stem+"gibbs_draws."+fmt('d',3,gibbs_count)('0')+".dat";
    fout.open(foutname.c_str());
    if (!fout) error("Error, usrvar, could not open " + foutname);

    fout << lrho << '\n';
    fout << gibbs_draws.size() << '\n';

    std::streamsize old_precision = fout.precision(REAL_DIG+1);
    fout.setf(std::ios::scientific,std::ios::floatfield);

    for (vector<realmat>::size_type j=0; j<gibbs_draws.size(); ++j) {
      for (INTEGER i=1; i<=lrho; ++i) {
        fout << gibbs_draws[j][i] << '\n';
      }
    }

    fout.precision(old_precision);
    fout.setf(std::ios::fmtflags(0),std::ios::floatfield);

    ++gibbs_count;

    gibbs_draws.clear();

    fout.clear(); fout.close();

    if ((*smooth)[0].size() == 0) return;

    m->set_parms(rho);

    INTEGER n = data.ncol();

    realmat mean(1,n+1,0.0);
    for (INTEGER i=0; i<N; ++i) {
      mean += (*smooth)[i];
    }
    mean = mean/N;
  
    realmat sdev(1,n+1,0.0);
    for (INTEGER i=0; i<N; ++i) {
      realmat z = (*smooth)[i] - mean;
      for (INTEGER t=1; t<=n+1; ++t) sdev[t] += z[t]*z[t];
    }
    for (INTEGER t=1; t<=n+1; ++t) sdev[t] = sqrt(sdev[t]/REAL(N-1));
    
    realmat smooth_mean = mean;
    realmat smooth_sdev = sdev;
    
    fill(mean,0.0);
    for (INTEGER i=0; i<N; ++i) {
      mean += (*filter)[i];
    }
    mean = mean/N;
  
    fill(sdev,0.0);
    for (INTEGER i=0; i<N; ++i) {
      realmat z = (*filter)[i] - mean;
      for (INTEGER t=1; t<=n+1; ++t) sdev[t] += z[t]*z[t];
    }
    for (INTEGER t=1; t<=n+1; ++t) sdev[t] = sqrt(sdev[t]/REAL(N-1));
    
    realmat filter_mean = mean;
    realmat filter_sdev = sdev;

    realmat raw;
    vecread("svsim.dat", raw, 2, data.ncol());

    foutname = stem + "filter.csv";
    fout.open(foutname.c_str());
    if (!fout) error("Error, smooth, cannot open fout");

    fout << "smooth_mean, smooth_sdev, filter_mean, filter_sdev, x, y" << '\n';
    for (INTEGER t=2; t<=n+1; ++t) {
      fout << smooth_mean[t] <<','<< smooth_sdev[t] <<','
        << filter_mean[t] <<','<< filter_sdev[t] <<','
        << raw(2,t-1) <<','<< raw(1,t-1) << '\n';
    }
    fout.clear(); fout.close(); 
  }
}
