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

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         afunc - Makes the coefficients a of an SNP density depend
              on x.

Syntax        #include "libsnp.h"

              class afunc {
                //...
              public:
                afunc(const std::vector<intvec>& ialpha, 
                  INTEGER imaxKz, INTEGER imaxIz, 
                  INTEGER iKx, INTEGER iIx, INTEGER ilx, INTEGER ilag);
                afunc(const afunc& iaf, const std::vector<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 ;
                INTEGER get_maxKz() const ;
                INTEGER get_maxIz() const;
                INTEGER get_Kx() const;
                INTEGER get_Ix() const;
                INTEGER get_lx() const;
                INTEGER get_lag() const;
                INTEGER get_nrowA() const;
                INTEGER get_ncolA() const;
                INTEGER get_nrow0() const;
                INTEGER get_nrowa0() const;
                const intvec& get_idx() const;
                const intvec& get_odx() const;
                const realmat& get_A() const;
                const realmat& get_a0() const;
                const std::vector<intvec>& get_alpha() const;
                const std::vector<intvec>& get_beta() const;
                void initialize_state();
                void set_A(const realmat& iA);
                void set_a0(const realmat& ia0);
                const realmat& operator() 
                  (const realmat& ix, realmat& dawa0, realmat& dawA);
                };

Declared in   libsnp.h

Description   Most methods are parameter input and output.  The main
              method is the application operator which computes afunc 
              and derivatives of a with respect to a0 and A.  The first
              element A[1] should be normalized to one; this is the
              user's responsibility; it is not enforced.  The normal
              indexing of A is as a vector: A[ij], ij=1,nrowA*ncolA,
              and the columns of dawA correspond to this indexing.
              The realmat a0 is a vector of intercept terms not in A
              and may be null.  Make sure nrow0>0 before accessing its
              elements.

Remarks       One first constructs an snpden; alpha of the constructor
              is as returned by method get_alpha of snpden.  maxKz and
              maxIz, if set smaller than Kz or Iz of the snpden, will
              eliminate the dependence of some rows of the coeffient 
              vector a on A.  These rows will then depend on a0. The
              index vector idx shows which depend on A and odx which 
              depend on a0.
              The second constructor copies the coefficients of the
              old afunc into the appropriate elements of the new.
              The only restriction is that ly implied by alpha 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.
              Implemented as Hermite polynomials for numerical stability.

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.
              INTEGER lu = f.get_ly();
              INTEGER lR = f.get_lR();
              INTEGER ltheta = lA+la0+lu+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;
                f.set_u(theta(seg(of+1,of+lu),1));
                of += lu;
                f.set_R(theta(seg(of+1,of+lR),1));
                REAL log_likelihood = 0.0;
                realmat dlfwa, dlfwu, dlfwR;
                realmat dawa0,dawA;
                realmat dllwa0;
                if (la0>0) realmat dllwa(1,la0,0.0);
                realmat dllwA(1,lA,0.0);
                realmat dllwu(1,ly,0.0);
                realmat dllwR(1,lR,0.0);
                for (INTEGER t=1; t<=sample_size; ++t) {
                  f.set_a(af(data("",t-1),dawa0,dawA));
                  log_likelihood += f.log_f(data("",t),dlfwa,dlfwu,dlfwR);
                  if ((la0>0) dllwa0 += dlfwa*dawa0; dllwA += dlfwa*dawA; 
                  dllwu += dlfwu; dllwR += dlfwR;
                }
                realmat dllwtheta;
                if (la0>0) dllwtheta = cbind(dllwA,dllwa0);
                else dllwtheta = dawA;
                dllwtheta = cbind(dllwtheta,dllwu);
                dllwtheta = cbind(dllwtheta,dllwR);
              //... end optimization loop

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

#include "libsnp.h"

using namespace scl;
using namespace std;

using namespace libsnp;

namespace {

  typedef pair<INTEGER,INTEGER> int_pair; 

  typedef pair<intvec,intvec> ivec_pair;

  struct ivec_pair_cmp {
    bool operator() (const ivec_pair& a, const ivec_pair& b) const
    {
      if (a.first.size()!=b.first.size() || a.second.size()!=b.second.size()) {
        error("Error, afunc, ivec_pair_cmp, length mismatch");
      }
      for (INTEGER i=1; i<=a.first.size(); ++i) {
        if (a.first[i] != b.first[i]) return (a.first[i] < b.first[i]);
      }
      for (INTEGER i=1; i<=a.second.size(); ++i) {
        if (a.second[i] != b.second[i]) return (a.second[i] < b.second[i]);
      }
      return false;
    }
  };
}

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

void libsnp::afunc::make_private()
{
  if (lx <= 0) {
    error("Error, afunc, afunc, lx must be positive");
  }

  if (lag < 0) {
    error("Error, afunc, afunc, lag cannot be negative");
  }

  P.get_multi(beta);
  ncolA = beta.size();
  INTEGER ly = alpha[0].size();

  intvec tmp_row0(la);
  intvec tmp_rowA(la);
  nrow0 = 0;
  nrowA = 0;
  for (INTEGER i=1; i<=la; ++i) {
    intvec multi = alpha[i-1];
    INTEGER count = 0;
    INTEGER sum = 0;
    for (INTEGER j=1; j<=ly; ++j) {
      if (multi[j] > 0) ++count;
      sum += multi[j];
    }
    switch (count) {
      case 0:                      //Intercept
        tmp_rowA[++nrowA] = i;
        break;
      case 1:                      //Main effect
        if (sum <= maxKz) {
          tmp_rowA[++nrowA] = i;
        }
        else {
          tmp_row0[++nrow0] = i;
        }
        break;
      default:                     //Interaction
        if (sum <= maxIz) {
          tmp_rowA[++nrowA] = i;
        }
        else {
          tmp_row0[++nrow0] = i;
        }
        break;
    }
  }

  if (nrow0>0) {
    odx.resize(nrow0);
    for (INTEGER i=1; i<=nrow0; ++i) {
      odx[i] = tmp_row0[i];
    }
  }

  idx.resize(nrowA);
  for (INTEGER i=1; i<=nrowA; ++i) {
    idx[i] = tmp_rowA[i];
  }

  intvec zero(lag>0?lx*lag:lx,0);
  if (beta[0] != zero) error("Error, afunc, afunc, beta must be ordered");

  if (nrow0>0) a0.resize(nrow0,1,0.0);
  A.resize(nrowA,ncolA,0.0);
  A[1]=1.0;
  a.resize(la,1,0.0);
}

libsnp::afunc::afunc(const vector<intvec>& ialpha, 
  INTEGER imaxKz, INTEGER imaxIz, INTEGER iKx, INTEGER iIx, 
  INTEGER ilx, INTEGER ilag)
: alpha(ialpha), la(ialpha.size()), maxKz(imaxKz), maxIz(imaxIz),
  Kx(iKx), Ix(iIx), lx(ilx), lag(ilag), x(lag>0?lx*lag:lx,1,0.0), 
  P('h',x,Kx,Ix)
{
  this->make_private();
}

libsnp::afunc::afunc(const afunc& iaf, const vector<intvec>& ialpha, 
  INTEGER imaxKz, INTEGER imaxIz, INTEGER iKx, INTEGER iIx, 
  INTEGER ilx, INTEGER ilag)
: alpha(ialpha), la(ialpha.size()), maxKz(imaxKz), maxIz(imaxIz),
  Kx(iKx), Ix(iIx), lx(ilx), lag(ilag), x(lag>0?lx*lag:lx,1,0.0),
  P('h',x,Kx,Ix)
{
  this->make_private();

  vector<intvec> old_alpha = iaf.get_alpha();

  if ( old_alpha[0].size() != alpha[0].size() ) {
    error("Error, afunc, afunc, y lengths differ");
  }

  bool nested = true;

  INTEGER old_lx = iaf.get_lx();
  INTEGER old_nrow0 = iaf.get_nrow0();
  INTEGER old_nrowA = iaf.get_nrowA();
  INTEGER old_ncolA = iaf.get_ncolA();
  realmat old_a0 = iaf.get_a0();
  realmat old_A = iaf.get_A();
  intvec old_odx = iaf.get_odx();
  intvec old_idx = iaf.get_idx();

  //Copy elements of iaf.A("",1) and a0 to A, matching by elements 
  //of alpha

  //Indexing conventions:
  //  alpha[i] indexes i=0,la-1
  //  idx[i] indexes i=1,nrowA
  //  alpha[idx[i]-1] indexes i=1,nrowA
  //  odx[i] indexes 1=1,nrow0
  //  alpha[odx[i]-1] indexes i=1,nrow0

  map<intvec,int,intvec_cmp> new_a;

  for (INTEGER i=1; i<=la; ++i) {
      new_a[alpha[i-1]] = i;
  }  

  #if defined SUN_CC_COMPILER
    map<intvec,int,intvec_cmp>::iterator end_a = new_a.end();
  #else
    map<intvec,int,intvec_cmp>::const_iterator end_a = new_a.end();
  #endif

  if (old_nrow0 > 0) {
    for (INTEGER i=1; i<=old_nrow0; ++i) {
      intvec old_multi = old_alpha[old_odx[i]-1];
      if (new_a.find(old_multi) == end_a) {
        nested = false;
      }
      else {
        a[new_a[old_multi]] = old_a0[i];
      }
    }
  }

  for (INTEGER i=1; i<=old_nrowA; ++i) {
    intvec old_multi = old_alpha[old_idx[i]-1];
    if (new_a.find(old_multi) == end_a) {
      nested = false;
    }
    else {
      a[new_a[old_multi]] = old_A(i,1);
    }
  }

  vector<intvec> old_beta = iaf.get_beta();


  INTEGER change;

  if ( lag == 0 ) {
    change = lx - old_lx;
  }
  else {
    if (old_lx != lx ) {
      error("Error, afunc, afunc, x lengths differ");
    }
    INTEGER old_lag = iaf.get_lag();
    change = lx*(lag - old_lag);
  }

  if (change > 0 ) {  //Assumes new elements of x added at the end
    intvec zero(change,0);
    for (INTEGER j=0; j<old_ncolA; ++j) {
      old_beta[j] = ibind(old_beta[j],zero);
    }
  }

  if (change < 0 ) {  //Assumes old values of x deleted from the end
    nested = false;
    intvec multi(lx);
    for (INTEGER j=0; j<old_ncolA; ++j) {
      bool keep = true;
      for (INTEGER i=lx+1; i<=old_lx; ++i) {
        if (old_beta[j][i] > 0) keep = false;
      }
      for (INTEGER i=1; i<=lx; ++i) {
        multi[i] = (keep ? old_beta[j][i] : -1);
      }
      old_beta[j] = multi; 
    }
  }

  //Copy elements of iaf.A to A, matching by elements of alpha and beta

  //Indexing conventions:
  //  A(i,j) indexes i=1,nrowA
  //  alpha[i] indexes i=0,la-1
  //  beta[j] indexes j=0,ncolA-1
  //  idx[i] indexes i=1,nrowA
  //  alpha[idx[i]-1] indexes i=1,nrowA

  map<ivec_pair,int_pair,ivec_pair_cmp> new_model;

  for (INTEGER j=1; j<=ncolA; ++j) {
    for (INTEGER i=1; i<=nrowA; ++i) {
      int_pair ij(i,j);
      new_model[ ivec_pair( alpha[idx[i]-1], beta[j-1] ) ] = ij;
    }
  }  

  #if defined SUN_CC_COMPILER
    map<ivec_pair,int_pair,ivec_pair_cmp>::iterator itr = new_model.end();
  #else
    map<ivec_pair,int_pair,ivec_pair_cmp>::const_iterator itr = new_model.end();
  #endif

  for (INTEGER j=1; j<=old_ncolA; ++j) {
    for (INTEGER i=1; i<=old_nrowA; ++i) {
      ivec_pair old_ab(old_alpha[old_idx[i]-1], old_beta[j-1] );
      if (new_model.find(old_ab) == itr) {
        nested = false;
      }
      else {
        int_pair ij = new_model[old_ab];
        A(ij.first,ij.second) = old_A(i,j);
      }
    }
  }
   
  if (!nested) warn("Warning, afunc, afunc, models not nested");

  if (nrow0 > 0) {
    for (INTEGER i=1; i<=nrow0; ++i) a0[i] = a[odx[i]];
  }
  for (INTEGER i=1; i<=nrowA; ++i) A(i,1) = a[idx[i]];

  //A[1]=1.0;  is not enforced.  It is the user's responsability.

}

libsnp::afunc::afunc()
: alpha(vector<intvec>()), la(0), maxKz(0), maxIz(0), Kx(0), Ix(0), 
  lx(0), lag(0), x(realmat()), P(poly()), beta(vector<intvec>()), 
  nrowA(0), ncolA(0), nrow0(0), idx(intvec()), odx(intvec()), 
  A(realmat()), a0(realmat()), a(realmat())
{ }

libsnp::afunc::afunc(const afunc& iaf)
: alpha(iaf.alpha), la(iaf.la), maxKz(iaf.maxKz), maxIz(iaf.maxIz),
  Kx(iaf.Kx), Ix(iaf.Ix), lx(iaf.lx), lag(iaf.lag), x(iaf.x), P(iaf.P),
  beta(iaf.beta), nrowA(iaf.nrowA), ncolA(iaf.ncolA), nrow0(iaf.nrow0),
  idx(iaf.idx), odx(iaf.odx), A(iaf.A), a0(iaf.a0), a(iaf.a)
{ }

afunc& libsnp::afunc::operator=(const afunc& iaf)
{
  if (this != &iaf) {
    alpha=iaf.alpha; la=iaf.la; maxKz=iaf.maxKz; maxIz=iaf.maxIz; Kx=iaf.Kx;
    Ix=iaf.Ix; lx=iaf.lx; lag=iaf.lag; x=iaf.x; P=iaf.P; beta=iaf.beta;
    nrowA=iaf.nrowA; ncolA=iaf.ncolA; nrow0=iaf.nrow0; idx=iaf.idx;
    odx=iaf.odx; A=iaf.A; a0=iaf.a0; a=iaf.a;
  }
  return *this;
}

const realmat& libsnp::afunc::operator() 
  (const realmat& ix, realmat& dawa0, realmat& dawA)
{
  if (dawA.nrow()!=la||dawA.ncol()!=nrowA*ncolA) dawA.resize(la,nrowA*ncolA);

  fill(a);
  fill(dawA);

  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];
    }
  }

  P.set_x(x);
  realmat basis;
  P.get_basis(basis);

  realmat Ax = A*basis;

  for (INTEGER i=1; i<=nrowA; ++i) {
      a[idx[i]] = Ax[i];
    for (INTEGER j=1; j<=ncolA; ++j) {
      dawA(idx[i],-nrowA+i+nrowA*j)=basis[j];  
    }
  }

  if (nrow0 > 0) {
    if ((dawa0.nrow()!=la)||(dawa0.ncol()!=nrow0)) dawa0.resize(la,nrow0);
    fill(dawa0);
    for (INTEGER i=1; i<=nrow0; ++i) {
      a[odx[i]] = a0[i];
      dawa0(odx[i],i)=1.0;
    }
  }
  else {
    realmat null;
    dawa0 = null;
  }

  return a;
}

const realmat& libsnp::afunc::operator() (const realmat& ix)
{
  fill(a);

  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];
    }
  }

  P.set_x(x);
  realmat basis;
  P.get_basis(basis);

  realmat Ax = A*basis;

  for (INTEGER i=1; i<=nrowA; ++i) {
      a[idx[i]] = Ax[i];
  }

  if (nrow0 > 0) {
    for (INTEGER i=1; i<=nrow0; ++i) {
      a[odx[i]] = a0[i];
    }
  }

  return a;
}
