#ifndef __FILE_XSTAR_NORM_H_SEEN__
#define __FILE_XSTAR_NORM_H_SEEN__

#include "libscl.h"
#include "tree_sim.h"
#include "tree_mf.h"

class xstar_norm : public scl::nleqns_base {
private:
  scl::realmat data;
  const INTEGER n_obs;
  const INTEGER mf_lags;
  const INTEGER HAC_lags;
  scl::realmat scale;
  scl::realmat z;
  scl::realmat theta;
  scl::realmat s;
  INTEGER t;
  tree_moment_function tree_mf;
  scl::gmm tree_gmm;
  INTEGER d;
  scl::realmat m;
  scl::realmat W;
  REAL logdetW;
  INTEGER rankW;
  scl::realmat S;
  scl::realmat R;
  scl::realmat Z;
  scl::realmat I;
public:
  xstar_norm(const scl::realmat& dat, INTEGER n, INTEGER mfl, INTEGER hacl)
  : data(dat), n_obs(n), mf_lags(mfl), HAC_lags(hacl)
  {
    scl::gmm tmp_gmm(&tree_mf, mf_lags, &data, n_obs, HAC_lags);
    tmp_gmm.set_correct_W_for_mean(true);
    tree_gmm.set_warning_messages(true);
    tree_gmm = tmp_gmm;
    d = tree_mf.get_d();
    scale.resize(d,1,1.0);
    I.resize(d,d,0.0);
    for (INTEGER i=1; i<=d; ++i) I(i,i) = 1.0;
  }
  bool set_theta(const scl::realmat& parms) 
  {
    theta = parms;
    if (theta.size() == 2) return true; else return false;
  }
  bool set_data_z
    (const scl::realmat& dat, const scl::realmat& u, scl::realmat& v)
  { 
    bool rv = true;
    data = dat;
    INTEGER rows = data.nrow();
    z = u;
    for (INTEGER i=1; i<=d; ++i) z[i] *= scale[i];
    if (rows != 2 || data.ncol() != n_obs) {
      scl::error("Error, xstar_norm.set_data, dat wrong size");
      rv = false;
    }
    REAL beta = theta[1];
    const REAL e = exp(1.0);
    s.resize(d,1);
    s[1] = log(1.0 - tanh(z[1]/4.0)) - log(beta);
    s[2] = z[2]/(1.0 - e);
    s[3] = z[3]/(1.0 - e);
    v = s;
    INTEGER t = n_obs - 2;
    INTEGER base = rows*(t - 1);
    data[base + 1] = s[2];
    data[base + 2] = s[3];
    data[base + 3] = 1.0 - log(beta);
    data[base + 4] = 0.0;
    data[base + 5] = s[1];
    data[base + 6] = 0.0;
    z = u;
    return rv;
  } 
  void set_scale(const scl::realmat& scl) 
  { if (scl.size() != d) scl::error("Error, xnorm.set_t, bad scl"); scale=scl;}
  void set_t(INTEGER obs) 
  { t=obs; if (t < 1 || n_obs < t) scl::error("Error, xnorm.set_t, bad t"); }
  void set_x(const scl::realmat& x) 
  { 
    INTEGER xsize = x.size();
    INTEGER rows = data.nrow();
    INTEGER base = rows*(t - 1);
    if (base + xsize > (n_obs-3)*rows) scl::error("Error, xnorm.set_x, bad x");
    for (INTEGER i=1; i<=xsize; ++i) data[base + i] = x[i];
  }
  scl::realmat get_data() { return data; }
  INTEGER get_d() { return d; }
  INTEGER get_n() { return n_obs; }
  scl::realmat get_s_inverse()
  {
    REAL beta = theta[1];
    const REAL e = exp(1.0);
    scl::realmat u(d,1);
    u[1] = 4.0*atanh( 1.0 - exp( s[1] + log(beta) ) );
    u[2] = (1.0 - e)*s[2];
    u[3] = (1.0 - e)*s[3];
    for (INTEGER i=1; i<=d; ++i) u[i] /= scale[i];
    return u;
  }
  bool get_f(const scl::realmat& x, scl::realmat& f)
  {
    bool rv = true;
    set_x(x);
    if (!tree_gmm.set_data(&data)) rv = false;
    tree_gmm(theta,m,W,logdetW,rankW,S);
    if (tree_gmm.get_W_numerr() > 0) {
      scl::warn("Warning, xstar_norm, numerr > 0");
      rv = false;
      f.resize(1,1,REAL_MAX);
      return rv;
    }
    if (cholesky(W,R) != d) rv = false;
    scl::realmat ztst = sqrt(n_obs)*R*m;
    REAL sum = 0.0;
    for (INTEGER i=1; i<=d; ++i) sum += pow(ztst[i]-z[i],2);
    sum = sqrt(sum);
    f.resize(1,1,sum);
//std::cerr << "m, zscale*z, S" << m << zscale*z << S;
    return rv;
  }
  bool get_F(const scl::realmat& x, scl::realmat& f, scl::realmat& F)
  {
    if (this->get_f(x,f)) {
      return nleqns_base::df(x,F);
    }
    else {
      return false;
    }
  }
};

#endif
