// K M Briggs 2009-10-05 11:30 

/* ode.h KMB 97 July 15 (was tests/tryode.cc)
ODE solver with AD computation of derivatives w.r.t. IVs and parameters.

NB: job=1 is not working! It now defaults to job=2
To do:
  job=1: dy/dp only

Notes:
  Integration may be continued without resetting Phi and Psi
Needs:
  ad.h from math++ (http://www-epidem.plantsci.cam.ac.uk/~kbriggs/math++1.9/ad.h)
  dopri5a.h and dopri5a.cc (my custom versions of Hairer's dopri5, included)
  newmat09 library (http://nz.com/webnz/robert/nzc_nm09.html)
Usage:
  #include "ode.h"
  void odes(int n, double x, ad y[], ad f[], ad p[]) {
    // ODE rhs functions
    f[0]  =(arbitrary expression using x,y,p);
    f[1]  =...;
    ...
    f[n-1]=...;
  }
  void ...() {
    const int n=(system state-space dimension), 
              m=(number of parameters in ODEs);
    double t_initial=..., t_final=...;
    RowVector y0(n),y(n),p(m);
    Matrix Psi(n,m), Phi(n,n);
    Psi=0.0; Phi=0.0; for (int i=1; i<=n; i++) Phi(i,i)=1.0;
    y0<<(IV);
    y=integrate(odes,t_initial,t_final,y0,p,Psi,Phi);
    // Now Phi=dy/dy0, Psi=dy/dp
  }
*/

#include <stdlib.h>
#include <math.h>
#include <iostream>
#include <iomanip>
#include <fstream>
#include "dopri5a.h" /* my custom version of dopri5 allowing client data */
#include "newmat.h"
#include "newmatio.h"
#include "newmatap.h"
#include "ad.h"

typedef Ad ad;
typedef void (*Ad_odes)(const int n, const double x, ad y[], ad f[], ad p[]);
ReturnMatrix integrate(Ad_odes odes, const double from, const double to, 
             RowVector& Y, RowVector& p, Matrix& Psi, Matrix& Phi, 
	     int job);

//int ad::dim; // for ad
Ad_odes p_odes; // pointer to ode rhs function

// main interface...
// assumes Psi and Phi are initialized to 0 and I.
ReturnMatrix integrate(Ad_odes odes, const double from, const double to, 
             RowVector& Y, RowVector& p, Matrix& Psi, Matrix& Phi, 
	     int job=2) {
  int i,j,l,code,itoler,iout,meth,n=Y.Ncols(),m=p.Ncols(),neqn;
  double uround,safe,fac1,fac2,beta,hmax,h,*y,*rtoler,*atoler,*cd;
  void (*solout)(long, double, double, double*, unsigned, int*)=NULL;
  void full_odes(unsigned n, double x, double *y, double *f, double* p);
  if (Psi.Ncols()!=m || Psi.Nrows()!=n || Phi.Ncols()!=n || Phi.Nrows()!=n) { 
    cerr<<"Matrix dimension error in integrate!\n"; exit(1);
  }
  switch (job) {
    case 0:  { dim=0;   neqn=n; break; }
    case 1:  job=2; 
          // numbering problem with job=1 !{ ad::dim=m;   neqn=n+n*m; break; }
    case 2:  ;
    default: { dim=n+m; neqn=n+n*m+n*n; }
  }
  long nmax,nstiff;
  unsigned nrdens,licont=neqn,icont[licont];
  RowVector Z=Y;
  FILE* fileout=stderr;
  double *psi,*phi;
  p_odes=odes;
  psi=y+n; phi=y+n+n*m;
  y=     (double*)calloc(neqn,sizeof(double));
  rtoler=(double*)calloc(neqn,sizeof(double));
  atoler=(double*)calloc(neqn,sizeof(double));
  cd=    (double*)calloc(m+3, sizeof(double));
  meth=itoler=1; nmax=10000;
  iout=nstiff=nrdens=0;
  uround=safe=beta=0.0; fac1=0.2; fac2=10;
  hmax=to-from; h=hmax>1.0 ? 1.0: hmax;
  for (i=0; i<n;    i++) atoler[i]=rtoler[i]=1.0e-6; // soln
  for (i=n; i<neqn; i++) atoler[i]=rtoler[i]=1.0e-3; // derivatives
  // IVs...
  for (i=1; i<=n; i++) {
    y[i-1]=Y(i);
    if (job>0) for (j=1; j<=m; j++) y[n+    (i-1)*m+(j-1)]=Psi(i,j);
    if (job>1) for (j=1; j<=n; j++) y[n+n*m+(i-1)*n+(j-1)]=Phi(i,j);
  }
  cd[0]=m; cd[1]=n; cd[2]=job; 
  for (i=0; i<m; i++) cd[i+3]=p(i+1);
  if (0) { // debug...
    cerr<<"Integrating from "<<from<<" to "<<to<<endl;
    cerr<<"with Y="<<Y<<" and p="<<p<<"  y="<<endl;
    for (i=0; i<neqn; i++) cerr<<y[i]<<" "; cout<<endl;
    cerr<<"neqn="<<neqn<<" n="<<n<<" m="<<m<<" job="<<job<<endl;
    //cerr<<"Psi=\n"<<Psi; cerr<<"Phi=\n"<<Phi;
  }
  code=dopri5(neqn, full_odes, from, y, to, rtoler, atoler, itoler, solout,
       iout, fileout, uround, safe, fac1, fac2, beta, hmax, h, nmax,
       meth, nstiff, nrdens, icont, licont, cd);
  if (1!=code) {
    cerr<<"dopri5 returns "<<code<<" trying to integrate from "<<from<<
          " to "<<to<<endl; exit(1);
  }
  if (job>0) for (l=n, i=1; i<=n; i++) for (j=1; j<=m; j++) Psi(i,j)=y[l++];
  for (i=1; i<=n; i++) {
    Z(i)=y[i-1];
    if (job>1) for (j=1; j<=n; j++) Phi(i,j)=y[l++];
  }
  free(cd); free(atoler); free(rtoler); free(y); 
  Z.Release();
  return Z;
}

void full_odes(unsigned neqn, double x, double *y, double *f, double* p) {
  // neqn=actual full # of ODEs dopri5 is integrating
  // n=# of ODEs in original system, before variational eqns added
  // m=# of parameters, p[3],...,p[m+2]
  int m=int(p[0]), n=int(p[1]), job=int(p[2]); 
  int i,j,k,l=n;
  ad F[neqn],Y[neqn],P[m];
  double s, *psi=y+n, *phi=y+n+n*m;
  //cerr<<"full_odes: neqn="<<neqn<<"  m="<<m<<"  n="<<n<<"  job="<<job<<"  dim="<<ad::dim<<endl;
  for (i=0; i<m; i++) { P[i]=p[i+3]; if (job>0) P[i].is_iv(i+n+1); }
  for (i=0; i<neqn; i++) Y[i]=y[i];
  if (job>1) for (i=0; i<n; i++) Y[i].is_iv(i+1);
  // call real ODEs with AD...
  p_odes(n,x,Y,F,P);
  // dy/dt...
  for (i=0; i<neqn; i++) f[i]=F[i][0];
  if (0==job) return;
  // dpsi/dt=df/dy.psi+df/dp...
  for (i=0; i<n; i++) for (j=0; j<m; j++) {
    for (s=0.0, k=0; k<n; k++) s+=F[i][k+1]*psi[k*m+j];
    f[l++]=s+F[i][n+j+1]; // +df/dp
  }
  if (1==job) return;
  // dphi/dt=df/dy.phi...
  for (i=0; i<n; i++) for (j=0; j<n; j++) {
    for (s=0.0, k=0; k<n; k++) s+=F[i][k+1]*phi[k*n+j];
    f[l++]=s;
  }
  return;
}

// old version...
void full_odes_ok(unsigned neqn, double x, double *y, double *f, double* p) {
  // neqn=actual full # of ODEs for dopri5 to integrate
  // n=# of ODEs in original system, before variational eqns added
  // m=# of parameters, p[3],...,p[m+2]
  int m=int(p[0]), n=int(p[1]), job=int(p[2]); 
  int i,j,k,l=n;
  ad F[neqn],Y[n],P[m];
  double s, *psi=y+n, *phi=y+n+n*m;
  //cerr<<neqn<<"  "<<m<<"  "<<n<<"  "<<job<<endl;
  for (i=0; i<m; i++) { P[i]=p[i+3]; if (job>0) P[i].is_iv(i+n+1); }
  for (i=0; i<n; i++) { Y[i]=y[i];   if (job>1) Y[i].is_iv(i+1); }
  p_odes(neqn,x,Y,F,P);
  // dy/dt...
  for (i=0; i<neqn; i++) f[i]=F[i][0];
  if (0==job) return;
  // dpsi/dt=df/dy.psi+df/dp...
  for (i=0; i<n; i++) for (j=0; j<m; j++) {
    for (s=0.0, k=0; k<n; k++) s+=F[i][k+1]*psi[k*m+j];
    f[l++]=s+F[i][n+j+1]; // +df/dp
  }
  if (1==job) return;
  // dphi/dt=df/dy.phi...
  for (i=0; i<n; i++) for (j=0; j<n; j++) {
    for (s=0.0, k=0; k<n; k++) s+=F[i][k+1]*phi[k*n+j];
    f[l++]=s;
  }
  return;
}
