#include "nrutil.h"
#include <math.h>
#include <assert.h>

#define KMAXX 8
#define IMAXX (KMAXX+1)
#define BS_SAFE	1	/* Wayne's factor, to make y'=y(0..1) give correct # digits */
#define SAFE1 0.25
#define SAFE2 0.7
#define REDMAX 1e-5	/* max stepsize factor reduction */
#define REDMIN 0.7	/* min stepsize factor reduction */
#define TINY 1e-30	/* prevents div by 0 */
#define SCALMX 0.1	/* 1/SCALMX is the max factor a stepsize can be inc */

/*
** From Press et al.
*/

/*
** ODE int wrapper function
** Return status: 0 = OK, 1 = step size too small, 2 = too many steps.
*/

#define MAXSTP 10000

#define KMAX 100
int kmax=KMAX, kount;
double xp[KMAX], yp[KMAX][2], dxsav;

int odeint(double ystart[], int nvar, double *x1, double x2, double eps,
    double h1, double hmin, int *nok, int *nbad,
    void(*derivs)(int,double,double[],double[]),
    int(*step)(double[],double[],int,double*,double,double,double[],
	double*,double*,void(*)(int,double,double[],double[])))
{
    int nstp, i, status;
    double xsav=0,x,hnext,hdid,h;
    double *yscal,*y,*dydx;

    eps *= BS_SAFE;
    yscal=dvector(0,nvar-1);
    y=dvector(0,nvar-1);
    dydx=dvector(0,nvar-1);
    x=*x1;
    h=SIGN(h1,x2-*x1);
    *nok=(*nbad)= kount = 0;
    for(i=0;i<nvar;i++) y[i]=ystart[i];
    if(kmax>0)xsav=x-dxsav*2;
    for(nstp=1;nstp<=MAXSTP;nstp++) {
	derivs(nvar,x,y,dydx);
	for(i=0;i<nvar;i++)
	    yscal[i]=fabs(y[i])+fabs(dydx[i]*h)+TINY;
	if(kmax>0 && kount <kmax-1 && fabs(x-xsav)>fabs(dxsav)) {
	    xp[++kount]=x;
	    for(i=0;i<nvar;i++) yp[kount][i]=y[i];
	    xsav=x;
	}
	if((x+h-x2)*(x+h-*x1)>0)h=x2-x;
	status = step(y,dydx,nvar,&x,h,eps,yscal,&hdid,&hnext,derivs);
	if(status != 0)
	    break;
	if(hdid == h) ++(*nok); else ++(*nbad);
	if((x-x2)*(x2-*x1)>=0.0) {
	    if(kmax) {
		xp[++kount]=x;
		for(i=0;i<nvar;i++) yp[kount][i]=y[i];
	    }
	    assert(status == 0);
	    break;	/* OK */
	}
	if(fabs(hnext) <= hmin)
	{
	    status = 1;
	    break;	/* nrerror("step size too small in odeint"); */
	}
	h=hnext;
	status = 2;	/* in the case this is the last loop iteration */
    }
    *x1=x;
    for(i=0;i<nvar;i++) ystart[i]=y[i];
    free_dvector(dydx,0,nvar-1);
    free_dvector(y,0,nvar-1);
    free_dvector(yscal,0,nvar-1);
    return status;
}


/*
** Modified Midpoint step.
*/
static void mmid(double y[], double dydx[], int nvar, double xs, double htot,
    int nstep, double yout[], void (*derivs)(int,double,double[],double[]))
{
    int n,i;
    double x,swap,h2,h,*ym,*Yn;
    ym=dvector(0,nvar-1);
    Yn=dvector(0,nvar-1);
    h=htot/nstep;
    for(i=0;i<nvar;i++) {
	ym[i] = y[i];
	Yn[i]=y[i]+h*dydx[i];
    }
    x=xs+h;
    derivs(nvar,x,Yn,yout);
    h2=2*h;
    for(n=2; n<=nstep; n++) {
	for(i=0;i<nvar;i++) {
	    swap=ym[i]+h2*yout[i];
	    ym[i]=Yn[i];
	    Yn[i]=swap;
	}
	x+=h;
	derivs(nvar,x,Yn,yout);
    }
    for(i=0;i<nvar;i++)
	yout[i]=0.5*(ym[i]+Yn[i]+h*yout[i]);
    free_dvector(Yn,0,nvar-1);
    free_dvector(ym,0,nvar-1);
}


static double **d, *x;

/*
** Polynomial extrapolation for B-S method.
*/
static void pzextr(int iest, double xest, double yest[], double yz[],
    double dy[], int nv)
{
    int k1,j;
    double q,f2,f1,delta,*c;

    c=dvector(0,nv-1);
    x[iest]=xest;
    for(j=0;j<nv;j++) dy[j]=yz[j]=yest[j];
    if(iest==1){
	for(j=0;j<nv;j++) d[j][1]=yest[j];
    } else {
	for(j=0;j<nv;j++) c[j]=yest[j];
	for(k1=1;k1<iest;k1++) {
	    delta=1/(x[iest-k1]-xest);
	    f1=xest*delta;
	    f2=x[iest-k1]*delta;
	    for(j=0;j<nv;j++) {
		q=d[j][k1];
		d[j][k1]=dy[j];
		delta=c[j]-q;
		dy[j]=f1*delta;
		c[j]=f2*delta;
		yz[j]+=dy[j];
	    }
	}
	for(j=0;j<nv;j++) d[j][iest]=dy[j];
    }
    free_dvector(c,0,nv-1);
}


/*
** (Optional) Rational function extrapolation for B-S method.
*/
static void rzextr(int iest, double xest, double yest[], double yz[],
    double dy[], int nv)
{
    int k,j;
    double yy,v,ddy=0,c,b1,b,*fx;

    fx=dvector(1,iest);
    x[iest]=xest;
    if(iest==1)
	for(j=0;j<nv;j++) {
	    yz[j]=yest[j];
	    d[j][1]=yest[j];
	    dy[j]=yest[j];
	}
    else {
	for(k=1;k<iest;k++)
	    fx[k+1]=x[iest-k]/xest;
	for(j=0;j<nv;j++) {
	    v=d[j][1];
	    d[j][1]=yy=c=yest[j];
	    for(k=2;k<=iest;k++) {
		b1=fx[k]*v;
		b=b1-c;
		if(b) {
		    b=(c-v)/b;
		    ddy=c*b;
		    c=b1*b;
		} else
		    ddy=v;
		if(k!=iest)v=d[j][k];
		d[j][k]=ddy;
		yy+=ddy;
	    }
	    dy[j]=ddy;
	    yz[j]=yy;
	}
    }
    free_dvector(fx,1,iest);
}




/*
** Burlish-Stoer integration.
** Return status: 0 = OK, 1 = step size too small.
*/
int bsstep(double y[], double dydx[], int nv, double *xx, double htry,
    double eps, double yscal[], double *hdid, double *hnext,
    void(*derivs)(int,double,double*,double*))
{
    void mmid(double Y[], double Dydx[], int nvar, double xs, double htot,
	int nstep, double yout[], void(*Derivs)(int,double,double[],double[]));
    int i,iq,k,kk,km=0;
    static int first=1,kmax,kopt;
    static double epsold=-1e30,xnew;
    double eps1,errmax=0,fact,h,red=0,scale=0,work,wrkmin,xest;
    double *err,*yerr,*ysav,*yseq;
    static double a[IMAXX+1];
    static double alf[KMAXX+1][KMAXX+1];
    static int nseq[IMAXX+1]={0,2,4,6,7,10,12,14,16,18};
    int reduct,exitflag=0;

    d=dmatrix(0,nv-1,1,KMAXX);
    err=dvector(1,KMAXX);
    x=dvector(1,KMAXX);
    yerr=dvector(0,nv-1);
    ysav=dvector(0,nv-1);
    yseq=dvector(0,nv-1);
    if(eps != epsold) {
	*hnext = xnew = -1e29;
	eps1=SAFE1*eps;
	a[1]=nseq[1]+1;
	for(k=1;k<=KMAXX;k++) a[k+1]=a[k]+nseq[k+1];
	for(iq=2;iq<=KMAXX;iq++) {
	    for(k=1;k<iq;k++)
		alf[k][iq]=pow(eps1,(a[k+1]-a[iq+1])/
		    ((a[iq+1]-a[1]+1.0)*(2*k+1)));
	}
	epsold=eps;
	for(kopt=2;kopt<KMAXX;kopt++)
	    if(a[kopt+1] > a[kopt]*alf[kopt-1][kopt]) break;
	kmax=kopt;
    }
    h=htry;
    for(i=0;i<nv;i++) ysav[i]=y[i];
    if(*xx != xnew || h != (*hnext)) {
	first =1;
	kopt=kmax;
    }
    reduct=0;
    while(1) {
	for(k=1;k<=kmax;k++) {
	    xnew=(*xx)+h;
	    if(xnew==(*xx)) return 1; /*nrerror("step size underflow in bsstep");*/
	    mmid(ysav,dydx,nv,*xx,h,nseq[k],yseq,derivs);
	    xest=SQR(h/nseq[k]);
	    rzextr(k,xest,yseq,y,yerr,nv);
	    if(k!=1) {
		errmax=TINY;
		for(i=0;i<nv;i++) errmax=FMAX(errmax,fabs(yerr[i]/yscal[i]));
		errmax/=eps;
		km=k-1;
		err[km]=pow(errmax/SAFE1,1.0/(2*km+1));
	    }
	    if(k!=1 && (k>=kopt-1 || first)) {
		if(errmax < 1.0) {
		    exitflag=1;
		    break;
		}
		if(k==kmax || k==kopt+1) {
		    red=SAFE2/err[km];
		    break;
		}
		else if (k==kopt && alf[kopt-1][kopt]<err[km]) {
			red=1/err[km];
			break;
		    }
		else if (kopt==kmax && alf[km][kmax-1] < err[km]) {
			red=alf[km][kmax-1]*SAFE2/err[km];
			break;
		    }
		else if (alf[km][kopt]<err[km]) {
		    red=alf[km][kopt-1]/err[km];
		    break;
		}
	    }
	}
	if(exitflag) break;
	red=FMIN(red,REDMIN);
	red=FMAX(red,REDMAX);
	h *= red;
	reduct=1;
    }
    *xx=xnew;
    *hdid=h;
    first=0;
    wrkmin=1e35;
    for(kk=1;kk<=km;kk++) {
	fact=FMAX(err[kk],SCALMX);
	work=fact*a[kk+1];
	if(work < wrkmin) {
	    scale=fact;
	    wrkmin=work;
	    kopt=kk+1;
	}
    }
    *hnext=h/scale;
    if(kopt >= k && kopt != kmax && !reduct) {
	fact=FMAX(scale/alf[kopt-1][kopt],SCALMX);
	if(a[kopt+1]*fact<=wrkmin) {
	    *hnext=h/fact;
	    kopt++;
	}
    }
    free_dvector(yseq,0,nv-1);
    free_dvector(ysav,0,nv-1);
    free_dvector(yerr,0,nv-1);
    free_dvector(x,1,KMAXX);
    free_dvector(err,1,KMAXX);
    free_dmatrix(d,0,nv-1,1,KMAXX);
    return 0;	/* OK */
}


#if 0	/* TEST */
/*
** Below here I only want to access things as 0 .. n-1
*/

void f(int n, double t, double y[n], double ydot[n])
{
    int i;
    for(i=0;i<n;i++)
	ydot[i] = y[i];
}

int main(void)
{
    double y = 1., x1=0, x2=100, eps=1e-16, h0=1, hmin=1e-12;
    int n=1, nok,nbad;

    kmax = 0;
    dxsav=0.1;
    printf("y(%.16g) = %.16g\n", x1, y);
    odeint(&y, n, &x1, x2, eps, h0, hmin, &nok, &nbad, f, bsstep);
    printf("Integration did %sfinish\n", x2==x1 ? "" : "NOT ");
    printf("y(%.16g) = %.16g\n", x1, y);
    printf("nok = %d, nbad = %d\n", nok, nbad);
    return 0;
}
#endif
