#include "RungeKutta.h"

#define ABSMAX(A,B) (fabs(A)>fabs(B)?fabs(A):fabs(B))
#define ABSMIN(A,B) (fabs(A)<fabs(B)?fabs(A):fabs(B))

void rkad( rkfloat xstart, rkfloat xend, 
	   int n, rkfloat ystart[], rkfloat yend[],
           rkfloat precision,
           void (*function)(rkfloat, rkfloat*, rkfloat*) ) {
  static rkfloat stepdone=FLT_MAX, stepnext=FLT_MAX;
  rkfloat *ytemp, *yscal, steptry;
  rkfloat x=xstart;
  int i;
  
  ytemp=(rkfloat*)malloc(n*sizeof(rkfloat));
  for (i=0;i<n;i++) ytemp[i]=ystart[i];
  yscal=(rkfloat*)malloc(n*sizeof(rkfloat));
  while (x<xend) {
    steptry = ABSMIN(xend-x,stepnext);
    for (i=0;i<n;i++) yscal[i]=ABSMAX(1.,ytemp[i]);
    rkqs(&x,n,ytemp,yscal,precision,steptry,&stepdone,&stepnext,function);
  }
  if (x>xend) rkck(x,n,ytemp,xend-x,yend,yscal,function);
  else  for (i=0;i<n;i++) yend[i]=ytemp[i];
  free(yscal);
  free(ytemp);
  return;
}

#define SAFETY   0.9
#define EXPUP   -0.20
#define EXPDOWN -0.25
#define MAXINCREMENT  5.
#define MAXDECREMENT   .1
#define MINSTEP  1.E-10

void rkqs( rkfloat *x, int n, rkfloat y[],
           rkfloat errscal[], rkfloat precision,
           rkfloat steptry, rkfloat *stepdone, rkfloat *stepnext,
           void (*function)(rkfloat, rkfloat*, rkfloat*) ) {
  rkfloat *ytemp, *yerr;
  rkfloat errmax;
  int i;

  ytemp=(rkfloat*)malloc(n*sizeof(rkfloat));
  yerr=(rkfloat*)malloc(n*sizeof(rkfloat));
  while (1) {
    rkck(*x,n,y,steptry,ytemp,yerr,function);
    *stepdone=steptry;
    errmax=0.;
    for (i=0;i<n;i++) errmax=ABSMAX(errmax,yerr[i]/errscal[i]);
    errmax/=(SAFETY*precision);
    if (errmax>1./SAFETY) {
      if (fabs(steptry)<MINSTEP) break;      
      *stepnext=steptry*ABSMAX(pow(errmax,EXPDOWN),MAXDECREMENT);
    } else if (errmax<SAFETY) {
      *stepnext=steptry*ABSMIN(pow(errmax,EXPUP),MAXINCREMENT);
      break;      
    } else break;
    steptry=*stepnext;
  }
  *x+=*stepdone;
  for (i=0;i<n;i++) y[i]=ytemp[i];
  free(ytemp);
  free(yerr);
  return;
}

#define RKEVAL 6

void rkck( rkfloat x, int n, rkfloat y[], 
           rkfloat h,        rkfloat ynext[], rkfloat yerr[],
           void (*function)(rkfloat, rkfloat*, rkfloat*) ) {

  static rkfloat a[RKEVAL] =
    {0.,0.2,0.3,0.6,1.,0.875};
  static rkfloat c[RKEVAL] =
    {37./378.,0.,250./621.,125./594,0.,512./1771.};
  static rkfloat cs[RKEVAL]=
    {2825./27648.,0.,18575./48384.,13525./55296.,277./14336.,0.25};
  static rkfloat b[(RKEVAL*(RKEVAL-1)/2)]=
    {.2,
     .075,0.225,
     .3,-.9,1.2,
     -11./54.,2.5,-10./27.,35./27.,
     1631./55296.,175./512.,575./13824.,44275./110592.,253./4096.};

  rkfloat *eval[RKEVAL+1], *ytemp;
  int i,j,k;

  for (i=0;i<RKEVAL+1;i++) eval[i]=(rkfloat*)malloc(n*sizeof(rkfloat));
  ytemp  =eval[RKEVAL];
  for (i=0;i<RKEVAL;i++) {
    for (j=0;j<n;j++) {
      ytemp[j]=y[j];
      for (k=0;k<i;k++) ytemp[j]+=h*b[i*(i-1)/2+k]*eval[i-1][j];
    }
    (*function)(x+a[i]*h,ytemp,eval[i]);
  }
  for (j=0;j<n;j++) {
    yerr[j]=ynext[j]=0.;
    for (i=0;i<RKEVAL;i++) {
      ynext[j]+=c[i]*eval[i][j];
      yerr[j]+=(c[i]-cs[i])*eval[i][j];
    } 
    ynext[j]=y[j]+h*ynext[j];
    yerr[j]*=h;
  }
  for (i=0;i<RKEVAL+1;i++) free(eval[i]);
}
