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

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.

-------------------------------------------------------------------------------

Class         ufunc - Makes the coefficients u of an SNP density be a
                      linear regression on x.

Syntax        #include "libsnp.h"

              class ufunc {
                //...
              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;
                INTEGER get_lx() const;
                INTEGER get_lag() const;
                bool is_intercept() const;
                bool is_regression() const;
                const realmat& get_b0() const;
                const realmat& get_B() const;
                void initialize_state();
                void set_b0(const realmat& ib0) 
                void set_B(const realmat& iB) 
                const realmat& operator() 
                  (const realmat& ix, realmat& duwb0, kronprd& duwB);
              };

Declared in   libsnp.h

Description   Most methods are parameter input and output.  The main
              method is the application operator which computes ufunc and 
              the derivatives of u with respect to b0 and B.  The normal
              indexing of B is as a vector: B[ij], ij=1,ly*lx and the
              columns of duwB correspond to this indexing.  The realmat
              b0 is a vector of intercept terms.  Either or both b0 and
              B can be null.  The methods is_intercept is true if b0 and
              duwb0 are not null and is_regression is true if B and duwB
              are not null.

Remarks       The input parameters ily and ilx set the dimension of B,
              which is ly by lx unless ilx=0, in which case it is null.
              If icept=false, then b0 is null, otherwise it is ly by 1.
              The second constructor copies the coefficients of the
              old ufunc into the appropriate elements of the new.
              The only restriction is that ly must be the same for both
              old and new.  If lag>0, the same is true for lx.  This 
              constructor is mainly for obtaining start values from a 
              previous fit.  

Reference     Gallant, A. Ronald, and George Tauchen (1992), "A
              Nonparametric Approach to Nonlinear Time Series Analysis:
              Estimation and Simulation," in Brillinger, David, Peter
              Caines, John Geweke, Emanuel Parzen, Murray Rosenblatt, and
              Murad S. Taqqu eds. (1992),  New Directions in Time Series
              Analysis, Part II.  Springer--Verlag, New York, 71-92.

Sample        #include "libsnp.h"
program       using namespace scl; 
              using namespace std;
              using namespace libsnp;
              //...
              INTEGER Kz=4; INTEGER Iz=2; INTEGER ly=3;
              INTEGER maxKz=2; INTEGER maxIz=0;
              INTEGER Kx=1; INTEGER Ix=0; INTEGER lx=3;
              INTEGER lag=0;
              snpden f(Kz,Iz,ly);
              afunc af(f.get_alpha(),maxKz,maxIz,Kx,Ix,lx,lag)
              INTEGER lA = af.get_nrowA*af.get_ncolA;
              INTEGER la0 = af.get_nrow0;  //Might be zero.
              ufunc uf(f.get_ly(),lx,lag,true);
              INTEGER lB = uf.get_ly()*uf_get_lx(); //Might be zero.
              INTEGER lb0 = 0;
              if (uf.is_intercept()) lb0=uf.get_ly();
              INTEGER lR = f.get_lR();
              INTEGER ltheta = lA+la0+lB+lb0+lR;
              realmat theta(ltheta,1);
              for (INTEGER i=1; i<=ltheta; ++i) theta[i] = start[i];
              theta[1]=1.0;  // Optimizer must hold theta[1] fixed.
              //... optimization loop
                INTEGER of=0;
                af.set_A(theta(seq(of+1,of+lA),1));
                of += lA;
                if (la0>0) af.set_a0(theta(seg(of+1,of+la0,1));
                of += la0;
                if (uf.is_intercept()) uf.set_b0(theta(seq(of+1,of+lb0),1));
                of += lb0;
                if (uf.is_regression()) uf.set_B(theta(seq(of+1,of+lB),1));
                of += lB;
                f.set_R(theta(seg(of+1,of+lR),1));
                REAL log_likelihood = 0.0;
                realmat dlfwa, dlfwu, dlfwR;
                realmat dawa0,dawA;
                realmat duwb0,duwB;
                realmat dllwa0;
                if (la0>0) realmat dllwa(1,la0,0.0);
                realmat dllwA(1,lA,0.0);
                realmat dllwb0;
                if (uf.is_intercept()) dllwb0.resize(1,ly,0.0);
                realmat dllwB;
                if (uf.is_regression()) dllwB.resize(1,ly*lx,0.0);
                realmat dllwA(1,lA,0.0);
                realmat dllwu(1,f.get_lu,0.0);
                realmat dllwR(1,f.get_lR,0.0);
                realmat dllwb0;
                if (uf.is_intercept()) dllwb0.resize(1,ly*ly,0.0);
                realmat dllwB;
                if (uf.is_regression()) dllwB.resize(1,ly*ly*lx,0.0);
                for (INTEGER t=2; t<=sample_size; ++t) {
                  f.set_a(af(data("",t-1),dawa0,dawA));
                  f.set_u(uf(data("",t-1),duwb0,duwB);
                  log_likelihood += f.log_f(data("",t),dlfwa,dlfwu,dlfwR);
                  if ((la0>0) dllwa0 += dlfwa*dawa0; dllwA += dlfwa*dawA;
                  if (uf.is_intercept()) dllwb0 += dlfwu*duwb0;
                  if (uf.is_regression()) dllwB += dlfwu*duwbB;
                  dllwR += dlfwR;
                }
                realmat dllwtheta;
                if (la0>0) dllwtheta = cbind(dllwA,dllwa0);
                else dllwtheta = dllwA;
                if (uf.is_intercept()) cbind(dllwtheta,dllwb0);
                if (uf.is_regression()) cbind(dllwtheta,dllwB);
                dllwtheta = cbind(dllwtheta,dllwR);
              //... end optimization loop


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

#include "libsnp.h"

using namespace scl;
using namespace std;

using namespace libsnp;

void libsnp::ufunc::initialize_state()
{
  for (INTEGER i=1; i<=lx*lag; ++i) {
    x[i] = 0.0;
  }
}

libsnp::ufunc::ufunc(INTEGER ily, INTEGER ilx, INTEGER ilag, bool icept)
  : ly(ily),u(ily,1,0.0),lx(ilx),lag(ilag),intercept(icept),regression(false)
{
  if (lag < 0 ) {
    error("Error, ufunc, ufunc, lag cannot be negative");
  }
  
  if (intercept) {
    b0.resize(ly,1,0.0);
  }

  if (lx > 0 ) {
    regression = true;
    x.resize(lx,1,0.0);
    if (lag > 0 ) {
      B.resize(ly,lx*lag,0.0);
      x.resize(lx*lag,1,0.0);
      initialize_state();
    } 
    else {
      B.resize(ly,lx,0.0);
      x.resize(lx,1,0.0);
    }
  }
}

libsnp::ufunc::ufunc(const ufunc& iuf, INTEGER ily, INTEGER ilx,
    INTEGER ilag, bool icept)
  : ly(ily),u(ily,1,0.0),lx(ilx),lag(ilag),intercept(icept),regression(false)
{
  if ( ly != iuf.get_ly() ) {
    error("Error, ufunc, ufunc, y lengths differ");
  }

  if (lag < 0 ) {
    error("Error, ufunc, ufunc, lag cannot be negative");
  }
  
  bool nested = true;

  if (intercept) {
    b0.resize(ly,1,0.0);
    if (iuf.is_intercept()) b0 = iuf.get_b0();
  }

  if (lx > 0 ) {

    regression = true;

    INTEGER old_lx = iuf.get_lx();

    if ( lag == 0 ) {

      B.resize(ly,lx,0.0);
      x.resize(lx,1,0.0);

      if (old_lx <= lx) { //Assumes new elements of x added at the end
        realmat old_B = iuf.get_B();
        for (INTEGER i=1; i<=ly*old_lx; ++i) B[i] = old_B[i];
  
        for (INTEGER i=1; i<=ly*old_lx; ++i) B[i] = old_B[i];
      }
      else if (old_lx > lx) { //Assumes old values of x deleted from the end
        nested = false;
        realmat old_B = iuf.get_B();
        for (INTEGER i=1; i<=ly*lx; ++i) B[i] = old_B[i];
      }
      else {
        error("Error, ufunc, ufunc, this should never happen");
      } 

    }
    else{

      if (old_lx != lx ) {
        error("Error, ufunc, ufunc, x lengths differ");
      }

      B.resize(ly,lx*lag,0.0);
      x.resize(lx*lag,1,0.0);
      initialize_state();

      INTEGER old_lag = iuf.get_lag();
      realmat old_B = iuf.get_B();

      INTEGER min_lag;
      if (lag >= old_lag) {
        min_lag = old_lag;
      }
      else {
        nested = false;
        min_lag = lag;
      }

      for (INTEGER k=1; k<=min_lag; ++k) {
        for (INTEGER j=1; j<=lx; ++j) {
          for (INTEGER i=1; i<=lx; ++i) { 
            B(i,lx*(k-1)+j)=old_B(i,lx*(k-1)+j);
          }
        }
      }

    }
  }

  if (iuf.intercept && !intercept) nested = false;
  if (iuf.regression && !regression) nested = false;

  if (!nested) warn("Warning, ufunc, ufunc, models not nested");
}

libsnp::ufunc::ufunc()
: ly(0), u(realmat()), lx(0), lag(0), x(realmat()), 
  intercept(false), regression(false), b0(realmat()), B(realmat())
{ }

libsnp::ufunc::ufunc(const ufunc& iuf)
: ly(iuf.ly), u(iuf.u), lx(iuf.lx), lag(iuf.lag), x(iuf.x), 
  intercept(iuf.intercept), regression(iuf.regression), b0(iuf.b0), B(iuf.B)
{ }

ufunc& libsnp::ufunc::operator=(const ufunc& iuf)
{
  if (this != &iuf) {
    ly=iuf.ly; u=iuf.u; lx=iuf.lx; lag=iuf.lag; x=iuf.x;
    intercept=iuf.intercept; regression=iuf.regression; b0=iuf.b0; B=iuf.B;
  }
  return *this;
}

const realmat& libsnp::ufunc::operator() 
  (const realmat& ix, realmat& duwb0, kronprd& duwB)
{

  fill(u);

  if (lag == 0 ) {
    for (INTEGER i=1; i<=lx; ++i) {
      x[i]=ix[i];
    }
  }
  else {
    for (INTEGER i=lx*lag; i>lx; --i) {
      x[i] = x[i-lx];
    }
    for (INTEGER i=1; i<=lx; ++i) {
      x[i] = ix[i];
    }
  }

  if (intercept) {

    u += b0;

    duwb0.resize(ly,ly,0.0);
    for (INTEGER i=1; i<=ly; ++i) {
      duwb0(i,i) = 1.0;
    }
  }
  else {
    realmat null;
    duwb0 = null;
  }
    
  if (regression) {

    u += B*x;

    realmat I(ly,ly,0.0);
    for (INTEGER i=1; i<=ly; ++i) {
      I(i,i) = 1.0;
    }
    duwB = kronprd(T(x),I);
  }
  else {
    kronprd null;
    duwB = null;
  }
  
  return u;
}

const realmat& libsnp::ufunc::operator() (const realmat& ix)
{
  fill(u);

  if (lag == 0 ) {
    for (INTEGER i=1; i<=lx; ++i) {
      x[i]=ix[i];
    }
  }
  else {
    for (INTEGER i=lx*lag; i>lx; --i) {
      x[i] = x[i-lx];
    }
    for (INTEGER i=1; i<=lx; ++i) {
      x[i] = ix[i];
    }
  }

  if (intercept) u += b0;
    
  if (regression) u += B*x;
  
  return u;
}

REAL libsnp::ufunc::stability() const
{
  INTEGER ier = 0;
  REAL rv = 0.0;

  if (lag == 0 || regression == false) return rv;

  if (lag == 1) {
   realmat companion = B;
   realmat V;
   rv = eigen(companion,V,ier);
  }
  else {
    realmat companion(lx*lag,lx*lag,0.0);
    for (INTEGER i=1; i<=lx*(lag-1); ++i) companion(lx+i,i) = 1.0;

    for (INTEGER j=1; j<=lx*lag; ++j) {
      for (INTEGER i=1; i<=lx; ++i) {
        companion(i,j) = B(i,j);
      }
    }
    realmat V;
    rv = eigen(companion,V,ier);
  }

  if (ier != 0) warn("Warning, ufunc, eigen value computation failed");
  return rv;
}






