#include "doubledouble.h"
#include "ddnrutil.h"
#include "ddbsstep.h"
#include <math.h>
#include <assert.h>

#define KMAXX 8
#define IMAXX (KMAXX+1)
#define BS_SAFE	doubledouble(1,0)	/* Wayne's factor, to make y'=y(0..1) give correct # digits */
#define SAFE1 doubledouble(0.25,0)
#define SAFE2 doubledouble(0.7,0)
#define REDMAX doubledouble(1e-5,0)	/* max stepsize factor reduction */
#define REDMIN doubledouble(0.7,0)	/* min stepsize factor reduction */
#define TINY doubledouble(1e-60,0)	/* prevents div by 0 */
#define SCALMX doubledouble(0.1,0)	/* 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 KMAX 100
int kmax=KMAX, kount;
doubledouble xp[KMAX], yp[KMAX][2], dxsav;

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

    eps *= BS_SAFE;
    yscal=ddvector(0,nvar-1);
    y=ddvector(0,nvar-1);
    dydx=ddvector(0,nvar-1);
    x=*x1;
    h=DDSIGN(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<=maxsteps;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)>doubledouble(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)>=doubledouble(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_ddvector(dydx,0,nvar-1);
    free_ddvector(y,0,nvar-1);
    free_ddvector(yscal,0,nvar-1);
    return status;
}


/*
** Modified Midpoint step.
*/
static void mmid(doubledouble y[], doubledouble dydx[], int nvar, doubledouble xs, doubledouble htot,
    int nstep, doubledouble yout[], void (*derivs)(int,doubledouble,doubledouble[],doubledouble[]))
{
    int n,i;
    doubledouble x,swap,h2,h,*ym,*Yn;
    ym=ddvector(0,nvar-1);
    Yn=ddvector(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]=doubledouble(0.5,0)*(ym[i]+Yn[i]+h*yout[i]);
    free_ddvector(Yn,0,nvar-1);
    free_ddvector(ym,0,nvar-1);
}


static doubledouble **d, *x;

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

    c=ddvector(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_ddvector(c,0,nv-1);
}


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

    fx=ddvector(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 != doubledouble(0)) {
		    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_ddvector(fx,1,iest);
}




/*
** Burlish-Stoer integration.
** Return status: 0 = OK, 1 = step size too small.
*/
int ddbsstep(doubledouble y[], doubledouble dydx[], int nv, doubledouble *xx, doubledouble htry,
    doubledouble eps, doubledouble yscal[], doubledouble *hdid, doubledouble *hnext,
    void(*derivs)(int,doubledouble,doubledouble*,doubledouble*))
{
    void mmid(doubledouble Y[], doubledouble Dydx[], int nvar, doubledouble xs, doubledouble htot,
	int nstep, doubledouble yout[], void(*Derivs)(int,doubledouble,doubledouble[],doubledouble[]));
    int i,iq,k,kk,km=0;
    static int first=1,kmax,kopt;
    static doubledouble epsold=-1e30,xnew;
    doubledouble eps1,errmax=0,fact,h,red=0,scale=0,work,wrkmin,xest;
    doubledouble *err,*yerr,*ysav,*yseq;
    static doubledouble a[IMAXX+1];
    static doubledouble 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=ddmatrix(0,nv-1,1,KMAXX);
    err=ddvector(1,KMAXX);
    x=ddvector(1,KMAXX);
    yerr=ddvector(0,nv-1);
    ysav=ddvector(0,nv-1);
    yseq=ddvector(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,doubledouble(1.0,0)/(2*km+1));
	    }
	    if(k!=1 && (k>=kopt-1 || first)) {
		if(errmax < doubledouble(1)) {
		    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_ddvector(yseq,0,nv-1);
    free_ddvector(ysav,0,nv-1);
    free_ddvector(yerr,0,nv-1);
    free_ddvector(x,1,KMAXX);
    free_ddvector(err,1,KMAXX);
    free_ddmatrix(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, doubledouble t, doubledouble *y, doubledouble *ydot)
{
    int i;
    for(i=0;i<n;i++)
	ydot[i] = y[i];
}

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

    kmax = 0;
    dxsav=0.1;
    cout << setprecision(30);
    cout << "y(" << x1 << ") = " << y << '\n';
    ddodeint(&y, n, &x1, x2, eps, h0, hmin, &nok, &nbad, f, ddbsstep);
    cout << "y(" << x2 << ") = " << y << '\n';
    cout << "nok = " << nok << ", nbad = " << nbad << '\n';
    return 0;
}
#endif
