// K M Briggs 2009-10-05 11:59 
// make bruss && ./bruss | colex 1 2 | p

// Hairer et al. p 201 example.
// KMB 98 June 29

#include "ode.h"

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

int dim;

int main() {
  int i,j,l;
  const int n=2, m=1;
  double B;
  double psifd[n][m],phifd[n][n];
  const double t0=0, t1=20;
  RowVector y(n),yout(n),p(m);
  Matrix Psi(n,m),Psifd(n,m);
  Matrix Phi(n,n),Phifd(n,n);
  for (B=2.88; B<3.08; B+=0.01) { // Hairer et al. p 201 example.
    p<<B;
    y<<1.3<<B;
    Phi=0.0; Psi=0.0; for (i=1; i<=n; i++) Phi(i,i)=1.0;
    y=integrate(odes,t0,t1,y,p,Psi,Phi,2);
    cout<<B<<" "<<Phi(1,1)+Psi(1,1)<<" "<<Phi(2,2)+Psi(2,1)<<endl;
  }
  return 0;
}

void odes(int n, double x, ad y[], ad f[], ad p[]) {
  ad t=sqr(y[0])*y[1];
  f[0]=1.0+t-(1.0+p[0])*y[0];
  f[1]=p[0]*y[0]-t;
  return;
}

// Another way, RK with AD...
void rk4(ad y[],double x,double h,ad yout[],ad p[]);
ReturnMatrix integrate1(Ad_odes odes, const double from, const double to,
             const RowVector& Y, const RowVector& P, Matrix& Psi, Matrix& Phi,
	     int job=2) {
  int i,j,n=Y.Ncols(),np=P.Ncols(),nsteps=100;
  double x=from, h=(to-from)/nsteps;
  dim=n+np;
  ad y[n],yout[n],p[np];
  //cerr<<"dim="<<n+np<<endl;
  for (i=0; i<n;  i++) { y[i]=Y(i+1); y[i].is_iv(i+1); }
  for (i=0; i<np; i++) { p[i]=P(i+1); p[i].is_iv(n+i+1); }
  for (i=0; i<nsteps; i++) {
    rk4(y,x,h,yout,p);
    x+=h;
    for (j=0; j<n; j++) y[j]=yout[j];
  }
  RowVector Z(n);
  for (i=0; i<n; i++) {
    Z(i+1)=yout[i][0];
    for (j=0; j<n; j++)  Phi(i+1,j+1)=yout[i][j+1];
    for (j=0; j<np; j++) Psi(i+1,j+1)=yout[i][n+j+1];
  }
  Z.Release();
  return Z;
}

void rk4(ad y[],double x,double h,ad yout[],ad p[]) {
  int i, n=3; 
  double xh,hh,h6;
  ad dym[n],dyt[n],yt[n],dydx[n];
  hh=h*0.5; h6=h/6.0; xh=x+hh; odes(n,x,y,dydx,p);
  for (i=0;i<n;i++) yt[i]=y[i]+hh*dydx[i];
  odes(n,xh,yt,dyt,p);  for (i=0;i<n;i++) yt[i]=y[i]+hh*dyt[i];
  odes(n,xh,yt,dym,p);  for (i=0;i<n;i++) { yt[i]=y[i]+h*dym[i]; dym[i] += dyt[i]; }
  odes(n,x+h,yt,dyt,p); for (i=0;i<n;i++) yout[i]=y[i]+h6*(dydx[i]+dyt[i]+2.0*dym[i]);
}
