#define HS_EXAMPLE
//#undef HS_EXAMPLE

#if defined HS_EXAMPLE

#include "libscl.h"

using namespace scl;
using namespace std;

// Hansen-Singleton (1984) Econometrica example.
// See Gallant, A. Ronald (1987) "Nolinear Statistical Models"
// Wiley, Chapter 5, for details.  Especially pp. 446 -- 451;

class hs_moment_function : public moment_function_base {
private:
  const realmat* data;
  INTEGER sample_size;
  INTEGER L;
  realmat theta;
public:
  hs_moment_function() : data(0), sample_size(0), L(0), theta() { }
  bool set_data(const realmat* datptr) {data = datptr; return true;}
  bool set_sample_size(INTEGER n) {sample_size = n; return true;}
  bool set_L(INTEGER lag) 
  {
    if (L != 0) warn("Warning, hs_moment_function, L = 0 regardless"); 
    L = 0; 
    return true;
  }
  bool set_theta(const realmat& parm) {theta = parm; return true;}
  INTEGER get_T0() {return 1;}
  INTEGER get_d() {return 3;}
  realmat operator() (INTEGER t) {realmat dm; return (*this)(t,dm); }
  realmat operator() (INTEGER t, realmat& dm)
  {
    realmat mt(3,1);
    dm.resize(3,2);

    REAL alpha = theta[1];
    REAL beta = theta[2];

    REAL y =  (*data)(1,t);
    REAL x =  (*data)(2,t);
    REAL z1 = (*data)(3,t);
    REAL z2 = (*data)(4,t);
    REAL z3 = (*data)(5,t);

    REAL e = beta * pow(y,alpha) * x - 1.0;

    REAL de1 = beta * log(y) * pow(y,alpha) * x;
    REAL de2 = pow(y,alpha) * x;

    mt[1] = z1 * e;
    mt[2] = z2 * e;
    mt[3] = z3 * e;

    dm(1,1) = z1 * de1;
    dm(1,2) = z1 * de2;
    dm(2,1) = z2 * de1;
    dm(2,2) = z2 * de2;
    dm(3,1) = z3 * de1;
    dm(3,2) = z3 * de2;

    return mt;
  }
};

class hs_nleqns : public nleqns_base {
private:
  gmm_base* g;
  realmat W;
public:
  hs_nleqns(gmm_base* hs_gmm, const realmat& Wmat) : g(hs_gmm), W(Wmat) { }
  bool set_W(const realmat& Wmat) {W = Wmat; return true;} 
  bool get_f(const realmat& theta, realmat& f)
  { 
     f.resize(1,1); 
     f[1] = (*g)(theta, W); 
     return true;
  }
  bool get_F(const realmat& theta, realmat& f, realmat& F)
  {
    get_f(theta,f); 
    F = ( T(g->get_m(theta)) * W ) * g->get_M(theta); 
    F += T(g->get_m(theta)) * ( T(W) * g->get_M(theta) );
    return true; 
  }
  /*
  bool get_F(const realmat& x, realmat& f, realmat& F)
  {
    if (this->get_f(x,f)) {
      return nleqns_base::df(x,F);
    }
    else {
      return false;
    }
  }
  */
};

int main(int argc, char** argp, char** envp)
{
  string filename = "ch6tbl1a.dat";
  ifstream fin;
  fin.open(filename.c_str());
  if (!fin) error("Error, cannot open " + filename);

  INTEGER t, mo;
  string yr;
  REAL nds_t, pop_t, nyse_t, def_t;

  realmat nds, pop, nyse, def;

  INTEGER n = 0;
  while (fin >> t >> yr >> mo >> nds_t >> pop_t >> nyse_t >> def_t) {
    nds.push_back(nds_t);
    pop.push_back(pop_t);
    nyse.push_back(nyse_t);
    def.push_back(def_t);
    ++n;
  }
  fin.close(); fin.clear();

  cout << starbox("/First and last three nds, pop, nyse, def//");
  cout << '\n';
  for (INTEGER t=1; t<=3; ++t) {
    cout << '\t' << fmt('f',5,1,nds[t]) << fmt('f',10,4,pop[t]) 
         << fmt('f',15,10,nyse[t]) << fmt('f',11,7,def[t]) << '\n';
  }
  cout << '\n';
  for (INTEGER t=n-2; t<=n; ++t) {
    cout << '\t' << fmt('f',5,1,nds[t]) << fmt('f',10,4,pop[t]) 
         << fmt('f',15,10,nyse[t]) << fmt('f',11,7,def[t]) << '\n';
  }

  n -= 2;

  realmat y(1,n), x(1,n), z(3,n);

  for (INTEGER t=1; t<=n; ++t) {
    y(1,t) = (nds[t+2]/pop[t+2])/(nds[t+1]/pop[t+1]);
    x(1,t) = (1.0 + nyse[t+2])*(def[t+1]/def[t+2]);
    z(1,t) = 1.0;
    z(2,t) = (nds[t+1]/pop[t+1])/(nds[t]/pop[t]);
    z(3,t) = (1.0 + nyse[t+1])*(def[t]/def[t+1]);
  }
  
  cout << starbox("/First and last three y, x, z//");
  cout << '\n';
  for (INTEGER t=1; t<=3; ++t) {
    cout << '\t' << fmt('f',9,5,y[t]) << fmt('f',9,5,x[t]) 
         << fmt('f',9,5,z(1,t)) << fmt('f',8,5,z(2,t)) 
	 << fmt('f',8,5,z(3,t)) <<'\n';
  }
  cout << '\n';
  for (INTEGER t=n-2; t<=n; ++t) {
    cout << '\t' << fmt('f',9,5,y[t]) << fmt('f',9,5,x[t]) 
         << fmt('f',9,5,z(1,t)) << fmt('f',8,5,z(2,t)) 
	 << fmt('f',8,5,z(3,t)) <<'\n';
  }

  realmat data(5,n);
  for (INTEGER t=1; t<=n; ++t) {
    data(1,t) = y[t];
    data(2,t) = x[t];
    data(3,t) = z(1,t);
    data(4,t) = z(2,t);
    data(5,t) = z(3,t);
  }

  cout << starbox("/First 6 columns of data//");
  cout << data("",seq(1,6));
  cout << starbox("/Last 6 columns of data//");
  cout << data("",seq(n-5,n));

  hs_moment_function hsmf;
  //INTEGER Lhac = 2;
  INTEGER Lhac = 0;
  gmm hs_gmm(&hsmf, 0, &data, n, Lhac);
  hs_gmm.set_correct_W_for_mean(false);
  hs_gmm.set_regularize_W(true);

  INTEGER d=hsmf.get_d();

  realmat W(d,d,0.0); 
  for (INTEGER t=1; t<=n; ++t) {
    W += z("",t)*T(z("",t));
  }
  W = invpsd(W);

  hs_nleqns hsnle(&hs_gmm,W);

  nlopt minimizer(hsnle);
  minimizer.set_lower_bound(0.0);
  minimizer.set_output(true, &cout);
  minimizer.set_check_derivatives(true);
  minimizer.set_warning_messages(true);

  realmat theta_start(2,1); 
  theta_start[1] = -0.4;
  theta_start[2] =  0.9;

  realmat theta_stop(2,1,0.0);

  if (minimizer.minimize(theta_start, theta_stop)) {
    cout << starbox("/First Stage Answer//") << theta_stop;
  }
  else {
    cout << starbox("/Failure/Check your derivatives!//");
    return 1;
  }

  realmat m;
  REAL logdetW;
  INTEGER rankW;

  theta_start = theta_stop;

  realmat S;
  hs_gmm(theta_start, m, W, logdetW, rankW, S);

  hsnle.set_W(W);

  cout << starbox("/Weighting Matrix W//") << W;
  cout << "\n\t\t\t logdetW = " << fmt('e',23,16,logdetW) << '\n';
  cout << "\n\t\t\t rankW = " << rankW << '\n';
  cout << "\n\t\t\t W_numerr = " << hs_gmm.get_W_numerr() << '\n';
  cout << starbox("/n*inv(W)//") << n*invpsd(W);
  cout << starbox("/n*S//") << n*S;

  //for (INTEGER i=1; i<=d; ++i) theta_start[i] += 0.001*ran(seed);

  if (minimizer.minimize(theta_start, theta_stop)) {
    cout << starbox("/Second Stage Answer//") << theta_stop;
  }
  else {
    cout << starbox("/Failure/Check your derivatives!//");
    return 1;
  }

  realmat theta = theta_stop;
  realmat Chat = invpsd(REAL(n)*(T(hs_gmm.get_M(theta))*W)*hs_gmm.get_M(theta));

  cout << starbox("/Chat//") << Chat;

  cout << starbox("n*s_n//") << "\n\t\t\t\t   " 
       << REAL(n)*hs_gmm(theta, m, W, logdetW, rankW, S) << '\n';

  realmat z_reflect;
  REAL ll = hs_gmm.likelihood(theta,z_reflect).log_den;

  cout << starbox("/z//") << "\n\t\t\t\t   " << z_reflect << '\n';
  cout << starbox("/log_likelihood//") << "\n\t\t\t\t   " << ll << '\n';

  return 0;
}

#else

#include "libscl.h"

using namespace scl;
using namespace std;

// Model: u_t = rho*u_{t-1} + e_t
//        x_t = cos(t) if t even; sin(t) else
//        y_t = theta[1]*y_{t-1} + theta[2]*x_t + u_t
// Data:  data(1,t) = y_t
//        data(2,t) = x_t

class eg_moment_function : public moment_function_base {
private:
  const realmat* dat;
  INTEGER sample_size;
  INTEGER L;
  realmat theta;
public:
  eg_moment_function() : dat(0), sample_size(0), L(0), theta() { }
  bool set_data(const realmat* datptr) {dat = datptr; return true;}
  bool set_sample_size(INTEGER n) {sample_size = n; return true;}
  bool set_L(INTEGER lag) {L = lag; return true;}
  bool set_theta(const realmat& parm) {theta = parm; return true;}
  INTEGER get_T0() {return 2*L+4;}
  INTEGER get_d() {return 2*L+2;}
  realmat operator() (INTEGER t) {realmat dm; return (*this)(t,dm); }
  realmat operator() (INTEGER t, realmat& dm)
  {
    realmat mt(2*L+2,1); dm.resize(2*L+2,2);
    for (INTEGER j=0; j<=2*L; j+=2) {
      REAL e = (*dat)(1,t-j)-theta[1]*(*dat)(1,t-j-1)-theta[2]*(*dat)(2,t-j);
      REAL de1 = -(*dat)(1,t-j-1);
      REAL de2 = -(*dat)(2,t-j);
      mt[j+1] = (*dat)(2,t-j)*e;
      mt[j+2] = (*dat)(2,t-j-1)*e;
      dm(j+1,1) = (*dat)(2,t-j)*de1;
      dm(j+1,2) = (*dat)(2,t-j)*de2;
      dm(j+2,1) = (*dat)(2,t-j-1)*de1;
      dm(j+2,2) = (*dat)(2,t-j-1)*de2;
    }
    return mt;
  }
};

class eg_nleqns : public nleqns_base {
private:
  gmm_base* g;
  realmat W;
public:
  eg_nleqns(gmm_base* eg_gmm, const realmat& Wmat) : g(eg_gmm), W(Wmat) {}
  bool set_W(const realmat& Wmat) {W = Wmat; return true;} 
  bool get_f(const realmat& theta, realmat& f)
  { f.resize(1,1); f[1] = (*g)(theta, W); return true;}
  bool get_F(const realmat& theta, realmat& f, realmat& F)
  { get_f(theta,f); F=2.0*(T(g->get_m(theta))*W)*g->get_M(theta); return true; }
};

int main(int argc, char** argp, char** envp)
{
  INTEGER n = 5000;
  INTEGER L = 1;
  INT_32BIT seed = 8913457;
  realmat theta(2,1,0.5);
  realmat data(2,n);
  REAL ulag = 0;  REAL ylag = 0;
  for (INTEGER t=1; t<=n+100; ++t) {
    REAL u = 0.5*ulag + 0.1*unsk(seed);
    REAL x = t%2 == 0 ? cos(t) : sin(t);
    REAL y = theta[1]*ylag + theta[2]*x + u;
    if (t>100) {data(1,t-100) = y; data(2,t-100) = x;}
    ulag = u; ylag = y;
  }
  eg_moment_function egmf;
  gmm eg_gmm(&egmf, L, &data, n, 1+INTEGER(pow(REAL(n),0.2)) );
  eg_gmm.set_regularize_W(true);
  INTEGER d = egmf.get_d();
  realmat W(d,d,0.0); for (INTEGER i=1; i<=d; ++i) W(i,i) = 1.0;
  eg_nleqns egnle(&eg_gmm, W);
  nlopt minimizer(egnle);
  minimizer.set_lower_bound(0.0); 
  minimizer.set_solution_tolerance(1.0e-3); // eg_gmm is quadratic
  minimizer.set_output(true, &cout);
  minimizer.set_check_derivatives(true);
  minimizer.set_warning_messages(true);
  realmat theta_start(2,1,0.0);
  if (minimizer.minimize(theta_start, theta)) {
    cout << starbox("/First Stage Answer//") << theta;
  } else { cout << starbox("/Failure//"); return 1; }
  realmat m;
  REAL logdetW;
  INTEGER rankW;
  theta_start = theta;
  realmat S;
  eg_gmm(theta_start, m, W, logdetW, rankW,S);
  egnle.set_W(W);
  cout << starbox("/Weighting Matrix W//") << W;
  cout << "\n\t\t\t logdetW = " << fmt('e',23,16,logdetW) << '\n';
  cout << "\n\t\t\t rankW = " << rankW << '\n';
  if (minimizer.minimize(theta_start, theta)) {
    cout << starbox("/Second Stage Answer//") << theta;
  } else { cout << starbox("/Failure//"); return 1; }
  realmat M = eg_gmm.get_M(theta);
  realmat Chat = invpsd(REAL(n)*(T(M)*W)*M);
  cout << starbox("/Chat//") << Chat;
  cout << starbox("z//") << hs_gmm.get_z(theta) << '\n';

  return 0;
}

#endif
