#ifndef __FILE_LIBSNP_H_SEEN__
#define __FILE_LIBSNP_H_SEEN__

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

Copyright (C) 2004, 2006.

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 "libscl.h"

namespace libsnp {

class snpden {
private:
  INTEGER Kz;                 //Degree of Hermite polynomial
  INTEGER Iz;                 //Degree of largest interaction
  INTEGER ly;                 //Dimension of y, z, and u, each ly by 1
  scl::realmat y;             //Random variable determined by SNP density
  scl::poly P;                //Multivariate Hermite polynomial, see libscl
  INTEGER la;                 //Dimension of Hermite coefficients a, la by 1
  scl::realmat a;             //Hermite polynomial coefficients, a is 1...la
  scl::realmat u;             //Location parameter, ly by 1
  INTEGER lR;                 //Dimension of R, lR=(ly*ly+ly)/2  
  scl::realmat R;             //Scale parameter, compressed upper triangular
  REAL e0;                    //Small epsilon to ensure positivity
  std::vector<scl::intvec> alpha;  //Multi indexes alpha is 0...la-1
  scl::realmat nmom;          //Even moments of the normal, grows as needed
  void make_private();        //Some common constructor code
public:
  snpden(INTEGER iKz, INTEGER iIz, INTEGER ily);
  snpden(const snpden& is, INTEGER iKz, INTEGER iIz, INTEGER ily);
  snpden();
  snpden(const snpden& is);
  ~snpden() { }
  snpden& operator=(const snpden& is);
  INTEGER get_Kz() const { return Kz; }
  INTEGER get_Iz() const { return Iz; }
  INTEGER get_ly() const { return ly; }
  INTEGER get_la() const { return la; }
  INTEGER get_lR() const { return lR; }
  const scl::realmat& get_a() const { return a; }
  const scl::realmat& get_u() const { return u; }
  const scl::realmat& get_R() const { return R; }
  REAL get_e0() const { return e0; }
  const std::vector<scl::intvec>& get_alpha() const { return alpha; }
  void get_Rgen(scl::realmat& Rgen) const;
  void set_a(const scl::realmat& ia) 
    {for (INTEGER i=1; i<=la; ++i) a[i]=ia[i];}
  void set_u(const scl::realmat& iu) 
    {for (INTEGER i=1; i<=ly; ++i) u[i]=iu[i];}
  void set_R(const scl::realmat& iR) 
    {for (INTEGER i=1; i<=lR; ++i) R[i]=iR[i];}
  void set_e0(REAL ie0) {e0 = ie0;}
  REAL log_f(const scl::realmat& iy); 
  REAL log_f(const scl::realmat& iy, scl::realmat& dlfwa, 
    scl::realmat& dlfwu, scl::realmat& dlfwR);
  void reg_poly(scl::realmat& a_reg) const; 
    //Returns reg poly coeff corresp to a.
  scl::realmat herm_poly(scl::realmat a_reg) const; 
    //Returns Herm coeff corresp to a_reg;
  scl::realmat ezl(const std::vector<scl::intvec>& lambda); 
    //lambda is 0..n-1, ezl 1..n
  void musig(scl::realmat& mu, scl::realmat& sig);
  void quad(INTEGER order, scl::realmat& abcissae, scl::realmat& weights);
    //make sure that order is large enough that weights sum to one
  scl::realmat sampy(INT_32BIT& seed) const;
};

class afunc {
private:
  std::vector<scl::intvec> alpha; //Multi idexes for z, alpha is 0...la-1
  INTEGER la;               //Dimension of Hermite coefficients a, la by 1
  INTEGER maxKz;            //Main effects of a > maxKz don't depend on x
  INTEGER maxIz;            //Interactions of a > maxIz don't depend on x
  INTEGER Kx;               //Degree of Hermite polynomial in x
  INTEGER Ix;               //Degree of largest interaction in x
  INTEGER lx;               //Dimension of x, x is lx by 1
  INTEGER lag;              //How many lags to put in x, 0 for cross section
  scl::realmat x;           //Regressors, either lx by 1 or lag*lx by 1
  scl::poly P;              //Multivariate Hermite polynomial, see libscl
  std::vector<scl::intvec> beta; //Multi indexes for x, beta is 0...ncolA-1
  INTEGER nrowA;            //Row dimension of A, A is nrowA by ncolA
  INTEGER ncolA;            //Col dimension of A and x's Hermite coefficients
  INTEGER nrow0;            //Row dimension of a0, may be 0, a0 is nrow0 by 1
  scl::intvec  idx;         //Gives mapping from rows of A to a, len is nrowA
  scl::intvec  odx;         //Gives mapping from rows of a0 to a, len is nrow0
  scl::realmat A;           //Coefficients, A[1]=1 required
  scl::realmat a0;          //Intercept terms that are not in A, may be null
  scl::realmat a;           //The scl::realmat returned by afunc, la by 1 
  void make_private();      //Some common constructor code
public:
  afunc(const std::vector<scl::intvec>& ialpha, 
    INTEGER imaxKz, INTEGER imaxIz, INTEGER iKx, INTEGER iIx, 
    INTEGER ilx, INTEGER ilag);
  afunc(const afunc& iaf, const std::vector<scl::intvec>& ialpha,
    INTEGER imaxKz, INTEGER imaxIz, INTEGER iKx, INTEGER iIx, 
    INTEGER ilx, INTEGER ilag);
  afunc();
  afunc(const afunc& iaf);
  ~afunc() { }
  afunc&  operator=(const afunc& iaf);
  INTEGER get_la() const { return la; }
  INTEGER get_maxKz() const { return maxKz; }
  INTEGER get_maxIz() const { return maxIz; }
  INTEGER get_Kx() const { return Kx; }
  INTEGER get_Ix() const { return Ix; }
  INTEGER get_lx() const { return lx; }
  INTEGER get_lag() const { return lag; }
  INTEGER get_nrow0() const { return nrow0; }   //Warning: nrow0 may be 0
  INTEGER get_nrowa0() const { return nrow0; }  //Warning: nrowa0 may be 0
  INTEGER get_nrowA() const { return nrowA; }
  INTEGER get_ncolA() const { return ncolA; }
  const scl::intvec& get_idx() const { return idx; }
  const scl::intvec& get_odx() const { return odx; } 
    //Warning: odx may be null
  const scl::realmat& get_a0() const { return a0; }
    //Warning: a0 may be null
  const scl::realmat& get_A() const { return A; }
  const std::vector<scl::intvec>& get_alpha() const { return alpha; }
  const std::vector<scl::intvec>& get_beta() const { return beta; }
  void initialize_state();
  void set_A(const scl::realmat& iA) 
    {for (INTEGER i=1; i<=nrowA*ncolA; ++i) A[i]=iA[i];}
  void set_a0(const scl::realmat& ia0) 
    {for (INTEGER i=1; i<=nrow0; ++i) a0[i]=ia0[i];}
  const scl::realmat& operator() (const scl::realmat& ix);
  const scl::realmat& operator() (const scl::realmat& ix, 
    scl::realmat& dawa0, scl::realmat& dawA);
};

class ufunc {
private:
  INTEGER ly;               //Dimension of y, z, and u, each ly by 1
  scl::realmat u;        //Location parameter, ly by 1
  INTEGER lx;               //Dim of input ix, ix is lx by 1, lx may be zero
  INTEGER lag;              //How many lags to put in x, 0 for cross section
  scl::realmat x;        //Regressors, either lx by 1 or lag*lx by 1
  bool intercept;           //True if intercept term in model
  bool regression;          //True if there is a regression, i.e. x not null
  scl::realmat b0;       //Intercept, b0 is either ly by 1 or null 
  scl::realmat B;        //Regr coef, B is ly by lx, ly by lx*lag, or null
public:
  ufunc(INTEGER ily, INTEGER ilx, INTEGER ilag, bool icept);
  ufunc(const ufunc& iuf, INTEGER ily, INTEGER ilx, INTEGER ilag, bool icept);
  ufunc();
  ufunc(const ufunc& iuf);
  ~ufunc() { }
  ufunc&  operator=(const ufunc& iuf);
  INTEGER get_ly() const { return ly; }
  INTEGER get_lx() const { return lx; }    //Warning: lx may be zero
  INTEGER get_lag() const { return lag; }  
  bool is_intercept() const { return intercept; }
  bool is_regression() const { return regression; }
  const scl::realmat& get_b0() const {return b0;} //Warning: b0 may be null
  const scl::realmat& get_B() const {return B;}   //Warning: B may be null
  void initialize_state();
  void set_b0(const scl::realmat& ib0) 
    {for (INTEGER i=1; i<=ly; ++i) b0[i]=ib0[i];}
  void set_B(const scl::realmat& iB) 
    {for (INTEGER i=1; i<=(lag>0?ly*lx*lag:ly*lx); ++i) B[i]=iB[i];}
  const scl::realmat& operator() (const scl::realmat& ix);
  const scl::realmat& operator() (const scl::realmat& ix, 
    scl::realmat& duwb0, scl::kronprd& duwB);
  REAL  stability() const;     //max singular value of the companion matrix
};

class rfunc {
private:
  INTEGER ly;                 //Dimension of y, z, and u, each ly by 1
  scl::realmat x;          //Input level effect vector
  INTEGER lx;                 //Dimension of x, which is lx by 1
  INTEGER squasher;           //Squash residuals: 0 none, 1 spline, 2 logistic
  REAL inflection;            //Inflection point of squasher
  INTEGER lR;                 //Dimension of R lR=(ly*ly+ly)/2.
  scl::realmat R;          //Scale parameter, R is lR by 1
  scl::realmat R0;         //Intercept term of GARCH, lR by 1
  INTEGER p;                  //Lags in MA part of GARCH, may be zero
  char Ptype;                 //Coded 's','d','f' for scalar, diagonal, full
  INTEGER rowsP;              //May be 0, 1, or ly
  INTEGER colsP;              //May be 0, p, or ly*p
  scl::realmat P;          //MA coefficient matrix, may be null
  INTEGER q;                  //Lags in AR part of GARCH, may be zero
  char Qtype;                 //Coded 's','d','f' for scalar, diagonal, full
  INTEGER rowsQ;              //May be 0, 1, or ly
  INTEGER colsQ;              //May be 0, q, or ly*q
  scl::realmat Q;          //AR coefficient matrix, may be null
  INTEGER v;                  //Lags in leverage effect of GARCH, may be zero
  char Vtype;                 //Coded 's','d','f' for scalar, diagonal, full
  INTEGER rowsV;              //May be 0, 1, or ly
  INTEGER colsV;              //May be 0, v, or ly*v
  scl::realmat V;          //Leverage effect coefficient matrix, may be null
  INTEGER w;                  //Lags in additive level effect, may be zero
  char Wtype;                 //Coded 's','d','f' for scalar, diagonal, full
  INTEGER rowsW;              //May be 0, 1, or ly
  INTEGER colsW;              //May be 0, w, or lx*w
  scl::realmat W;          //Level effect coefficient matrix, may be null
  INTEGER lRparms;            //Sum lengths of all parameter matricies
  mutable scl::realmat Rparms;  //Rparms=(vecR0,vecP,vecQ,vecV,vecW)
  struct sE_k {               //An element of the Pstate & Vstate  vector
    scl::realmat y;        //Observed y
    scl::realmat u;        //Predicted mean of y
    scl::realmat duwb0;    //Derivative of u wrt b0 
    scl::kronprd duwB;     //Derivative of u wrt B
    bool initial;             //True if initial conditions
    sE_k();
    sE_k(const scl::realmat&, const scl::realmat&, 
      const scl::realmat&, const scl::kronprd&, bool);
    sE_k(const sE_k&);
    sE_k& operator=(const sE_k&);
    ~sE_k() {}
  };
  struct sH_k {               //An element of the Qstate vector
    scl::realmat H;        //Variance, R is the Cholesky factor of H
    scl::realmat dHwb0;    //Derivative of H wrt b0 
    scl::realmat dHwB;     //Derivative of H wrt B
    scl::realmat dHwR0;    //Derivative of H wrt R0
    scl::realmat dHwP;     //Derivative of H wrt P
    scl::realmat dHwQ;     //Derivative of H wrt Q
    scl::realmat dHwV;     //Derivative of H wrt V
    scl::realmat dHwW;     //Derivative of H wrt W
    bool initial;             //True if initial conditions
    sH_k();
    sH_k(const scl::realmat&, const scl::realmat&, const scl::realmat&,
        const scl::realmat&, const scl::realmat&, const scl::realmat&, 
        const scl::realmat&, const scl::realmat&, bool);
    sH_k(const sH_k&);
    sH_k& operator=(const sH_k&);
    ~sH_k() {}
  };
  struct sX_k {               //An element of the Wstate vector
    scl::realmat x;        //Observed x
    bool initial;             //True if initial conditions
    sX_k();
    sX_k(const scl::realmat&, bool);
    sX_k(const sX_k&);
    sX_k& operator=(const sX_k&);
    ~sX_k() {}
  };
  std::vector<sE_k> Pstate;  //The state vector for the MA part
  std::vector<sH_k> Qstate;  //The state vector for the AR part
  std::vector<sE_k> Vstate;  //The state vector for the leverage effect
  std::vector<sX_k> Wstate;  //The state vector for the level effect
  bool factor_warn;
  void make_private();       //Some common constructor code
public:
  rfunc(INTEGER ily, INTEGER ilx, 
        INTEGER squasher, REAL inflection,
        INTEGER ip, char iPtype, INTEGER iq, char iQtype,
        INTEGER iv, char iVtype, INTEGER iw, char iWtype);
  rfunc(const rfunc& irf, INTEGER ily, INTEGER ilx, 
        INTEGER squasher, REAL inflection,
        INTEGER ip, char iPtype, INTEGER iq, char iQtype,
        INTEGER iv, char iVtype, INTEGER iw, char iWtype);
  rfunc();
  rfunc(const rfunc& irf);
  ~rfunc() { }
  rfunc&  operator=(const rfunc& irf);
  INTEGER get_ly() const { return ly; }
  INTEGER get_lx() const { return lx; }
  INTEGER get_lR() const { return lR; }
  const scl::realmat& get_R0() const { return R0; }
  INTEGER get_p() const { return p; }
  char get_Ptype() const { return Ptype; }
  INTEGER get_rowsP() const { return rowsP; }
  INTEGER get_colsP() const { return colsP; }
  const scl::realmat& get_P() const { return P; }
  INTEGER get_q() const { return q; }
  char get_Qtype() const { return Qtype; }
  INTEGER get_rowsQ() const { return rowsQ; }
  INTEGER get_colsQ() const { return colsQ; }
  const scl::realmat& get_Q() const { return Q; }
  INTEGER get_v() const { return v; }
  char get_Vtype() const { return Vtype; }
  INTEGER get_rowsV() const { return rowsV; }
  INTEGER get_colsV() const { return colsV; }
  const scl::realmat& get_V() const { return V; }
  INTEGER get_w() const { return w; }
  char get_Wtype() const { return Wtype; }
  INTEGER get_rowsW() const { return rowsW; }
  INTEGER get_colsW() const { return colsW; }
  const scl::realmat& get_W() const { return W; }
  INTEGER get_lRparms() const { return lRparms; }
  const scl::realmat& get_Rparms() const; //mutable Rparms updated
  void initialize_state();
  void set_R0(const scl::realmat& iR0) 
    {for (INTEGER i=1; i<=lR; ++i) R0[i]=iR0[i];}
  void set_P(const scl::realmat& iP) 
    {for (INTEGER i=1; i<=rowsP*colsP; ++i) P[i]=iP[i];}
  void set_Q(const scl::realmat& iQ) 
    {for (INTEGER i=1; i<=rowsQ*colsQ; ++i) Q[i]=iQ[i];}
  void set_V(const scl::realmat& iV) 
    {for (INTEGER i=1; i<=rowsV*colsV; ++i) V[i]=iV[i];}
  void set_W(const scl::realmat& iW) 
    {for (INTEGER i=1; i<=rowsW*colsW; ++i) W[i]=iW[i];}
  void set_Rparms(const scl::realmat& iRparms);
  const scl::realmat& operator() (
    const scl::realmat& iy, const scl::realmat& iu, 
    const scl::realmat& ix,
    const scl::realmat& iduwb0, const scl::kronprd& iduwB, 
    scl::realmat& dRwb0, scl::realmat& dRwB, 
    scl::realmat& dRwR0, scl::realmat& dRwP, scl::realmat& dRwQ, 
    scl::realmat& dRwV,  scl::realmat& dRwW);
  const scl::realmat& operator() (
    const scl::realmat& iy, const scl::realmat& iu, 
    const scl::realmat& ix,
    const scl::realmat& iduwb0, const scl::kronprd& iduwB, 
    scl::realmat& dRwb0, scl::realmat& dRwB, 
    scl::realmat& dRwRparms);
  const scl::realmat& operator() (
    const scl::realmat& iy, const scl::realmat& iu, 
    const scl::realmat& ix);
  REAL  stability() const;  //max singular value of P & Q companion matrix
};

}
#endif

