// RK.cc

#include <iostream.h>
#include <iomanip.h>
#include <fstream.h>
#include <math.h>
#include <stdlib.h>
#include "Nrutil.h"
#include "RKtest.h"

// ODEINT.C
    
#define MAXSTP 10000
    
void odeint(double *ystart,int nvar,double x1,double x2,double eps,
	    double h1,double hmin,int *nok,int *nbad,
	    void(*derivs)(const double, const int, double *, double *),
	    void(*rkqc)(double *, double *, int, double *, double, double,
	    double *, double *, double *, 
            void(*)(const double, const int, double *, double *)))
// odeint.c (from Numerical Recipes in C, 1st. ed., Ch. 15)
// Runge-Kutta driver with adaptive stepsize control.  Integrate starting 
// values ystart[1..nvar] from x1 to x2 with accuracy eps, storing intermediate
// results in global variables.  h1 should be set as a guessed first stepsize, 
// hmin as the minimum number of good and bad (but retried and fixed) steps
// taken, and ystart is replaced by values at the end of the integration
// interval.  derivs is the user-supplied routine for calculating the 
// right-hand side derivative, while rkqc is the name of the stepper routine
// to be used.
//
// Modifications:
// 26 Feb 1997:  Added code to output the step size after every step (hdid)
//               if the global variable statflag==1.  Position and step size
//               are written to a global variable "stepfilename" which is
//               defined in the calling routine.
// 23 Mar 1997:  Changed TINY from a parameter set by a preprocessor
//               directive to one set by a global variable "tiny".
// 13 Apr 1997:  Added code to display sum of amplitudes squared
//               at each step in stepsize file.               
{
  int nstp,i;
  double xsav,x,hnext,hdid,h;
  double *yscal,*y,*dydx;
    
  yscal=dvector(1,nvar);
  y=dvector(1,nvar);
  dydx=dvector(1,nvar);
  x=x1;
  h=(x2 > x1) ? fabs(h1) : -fabs(h1);
  *nok = (*nbad) = 0;
  for (i=1;i<=nvar;i++) y[i]=ystart[i];
  for (nstp=1;nstp<=MAXSTP;nstp++) {  // Take at most MAXSTP steps
    (*derivs)(x,nvar,y,dydx);
    for (i=1;i<=nvar;i++)             // Scaling used to monitor accuracy
      yscal[i]=fabs(y[i])+fabs(dydx[i]*h)+tiny;
    if ((x+h-x2)*(x+h-x1) > 0.0) {    // If step can overshoot end,
      hi=h;                           // cut down stepsize.
      h=x2-x;
    }
    (*rkqc)(y,dydx,nvar,&x,h,eps,yscal,&hdid,&hnext,derivs);
      if (statflag==1) {              // Output integration statistics
        ofstream StepFile;
        double ampsum=0;
        StepFile.open(stepfilename, ios::app);   // Open data file
        if (!StepFile) {                         // Error checking
          cerr << "File could not be opened" << endl;
        exit(1);      
        }
        for (i=1;i<=nvar;i++)
	  ampsum+=y[i]*y[i];
        StepFile.precision(6);              // Set output precision
        StepFile.unsetf(ios::fixed);        // Unset fixed notation
        StepFile.setf(ios::scientific);     // Set scientific notation
	StepFile << setw(8) << x << setw(15) << hdid;
        StepFile.unsetf(ios::scientific);   // Unset scientific notation
        StepFile.setf(ios::fixed);          // Set fixed notation
	StepFile << setprecision(10) << setw(17) << ampsum << endl;
        StepFile.close();
      }
      if (hdid == h) ++(*nok); else ++(*nbad);
      if ((x-x2)*(x2-x1) >= 0.0) {    // Are we done?
	for (i=1;i<=nvar;i++) ystart[i]=y[i];
	free_dvector(dydx,1,nvar);
	free_dvector(y,1,nvar);
	free_dvector(yscal,1,nvar);
	return;                       // Normal exit
      }
      if (fabs(hnext) <= hmin) nrerror("ODEINT1");
      h=hnext;
  }
  nrerror("ODEINT2");
}
    
#undef MAXSTP

// RKQC.C
    
#define PGROW -0.20
#define PSHRNK -0.25
#define FCOR 0.06666666		// 1.0/15.0
#define SAFETY 0.9
#define ERRCON 6.0e-4
    
void rkqc(double *y, double *dydx, int n, double *x, double htry,
	  double eps, double *yscal, double *hdid, double *hnext,
	  void (*derivs)(const double, const int, double *, double *))
// rkqc.c (from Numerical Recipes, 1st ed., Ch. 15)
// Fifth-order Runge-Kutta step with monitoring of local truncation error
// to ensure accuracy and adjust stepsize.  Input are the dependent variable
// vector y[1..n] and its derivative dydx[1..n] at the starting value of the 
// independent variable x.  Also input are the stepsize to be attempted htry,
// the required accuracy eps, the vector yscal[1..n] against which the error
// is scaled.  On output, y and x are replaced by their new values, hdid is 
// the stepsize that was actually accomplished, and hnext is the estimated
// next stepsize.  derivs is the user-supplied routine that computes the
// right-hand side derivatives.
//
// Modifications:
//     11 Apr 1997:  Added integration diagnostics code.  Should probably
//              be commented out once program is running correctly
//              since this routine is called many times!

{
  int i;
  double xsav,hh,h,temp,errmax;
  double *dysav,*ysav,*ytemp;
    
  cout.setf(ios::scientific);      // Diagnostics code
  cout.precision(4);
  int OnPrintList;
  
  dysav=dvector(1,n);
  ysav=dvector(1,n);
  ytemp=dvector(1,n);
  xsav=(*x);                       // Save initial values
  for (i=1;i<=n;i++) {
    ysav[i]=y[i];
    dysav[i]=dydx[i];
  }
  h=htry;                          // Set stepsize to initial trial value 
  for (;;) {
    hh=0.5*h;                      // Take two half steps
    rk4(ysav,dysav,n,xsav,hh,ytemp,derivs);
    *x=xsav+hh;
    (*derivs)(*x,n,ytemp,dydx);
    rk4(ytemp,dydx,n,*x,hh,y,derivs);
    *x=xsav+h;
    if (*x == xsav) nrerror("RKQC");
    rk4(ysav,dysav,n,xsav,h,ytemp,derivs);  // Take the large step
    errmax=0.0;                    // Evaluate accuracy
    system("clear");               // Diagnostics code
    for (i=1;i<=n;i++) {
      ytemp[i]=y[i]-ytemp[i];      // ytemp now contains the error estimate
      temp=fabs(ytemp[i]/yscal[i]);
      OnPrintList = (i%10==9 && i<=100) || ((i%100)==9) || (i%10==9 && i>=(n-100));
      if ( DiagnosticsOn && OnPrintList ) { 
	cout << "x= "       << setw(6)  << xsav 
             << " ["        << setw(3)  << i << "]"
             << " y= "      << setw(9)  << y[i]
             << " ytemp= "  << setw(9)  << ytemp[i]
             << " yscal= "  << setw(9)  << yscal[i]
             << " temp = "  << setw(9)  << temp 
             << " errmax= " << setw(9)  << errmax << endl;
      }     
      if (errmax < temp) errmax=temp;
    }
    errmax /= eps;                 // Scale relative to required tolerance
    if (errmax <= 1.0) {           // Step succeeded.  Compute next stepsize
      *hdid=h;
      *hnext=(errmax > ERRCON ?
	      SAFETY*h*exp(PGROW*log(errmax)) : 4.0*h);
      break;
    }
    h=SAFETY*h*exp(PSHRNK*log(errmax));  // Truncation error too large,
                                         //   reduce stepsize
  }
  for (i=1;i<=n;i++) y[i] += ytemp[i]*FCOR;  // Mop up 5th order trunc. error
  //  for (i=1;i<=1000;i++){};
  free_dvector(ytemp,1,n);
  free_dvector(dysav,1,n);
  free_dvector(ysav,1,n);
}
    
#undef PGROW
#undef PSHRNK
#undef FCOR
#undef SAFETY
#undef ERRCON


// RK4.C

void rk4(double *y, double *dydx, int n, double x, double h, double *yout,
	 void (*derivs)(const double, const int, double *, double *))
// rk4.c (from Numerical Recipes, 1st ed., Ch. 15)
// Given values for variables y[1..n] and their derivatives dydx[1..n] known
// at x, use the fourth-order Runge-Kutta method to advance the solution over
// an interval h and return the incremented variables as yout[1..n] which
// need not be a distinct array from y.  The user supplies the routine 
// derivs(x,n,y,dydx) which returns derivatives dydx at x.
{
  int i;
  double xh,hh,h6,*dym,*dyt,*yt;
    
  dym=dvector(1,n);
  dyt=dvector(1,n);
  yt=dvector(1,n);
  hh=h*0.5;
  h6=h/6.0;
  xh=x+hh;
  for (i=1;i<=n;i++) yt[i]=y[i]+hh*dydx[i];    // First step
  (*derivs)(xh,n,yt,dyt);                      // Second step
  for (i=1;i<=n;i++) yt[i]=y[i]+hh*dyt[i];   
  (*derivs)(xh,n,yt,dym);                      // Third step
  for (i=1;i<=n;i++) {
    yt[i]=y[i]+h*dym[i];
    dym[i] += dyt[i];
  }
  (*derivs)(x+h,n,yt,dyt);                     // Fourth step
    for (i=1;i<=n;i++)    // Accumulate in increments with proper weights
      yout[i]=y[i]+h6*(dydx[i]+dyt[i]+2.0*dym[i]);
    free_dvector(yt,1,n);
    free_dvector(dyt,1,n);
    free_dvector(dym,1,n);
}

void rkfixed(double *ystart,int nvar,double x1,double x2,double dx,
	     void(*derivs)(const double, const int, double *, double *),
             const float *nvec, float *normdB, float *normdBold)
// rkfixed.c (loosely based on rkdumb.c in Numerical Recipes in C, 
//            1st. ed., Ch. 15)
// Starting from initial values y[1..nvar] for nvar functions known
// at x1, use 4th order Runge-Kutta to advance the solution in increments
// of dx to x2.  If the final step will cause the solution to advance
// beyond x2, then the final step size will be adjusted so that it does not.
// the user-supplied routine derivs(x,v,dvdx) evaluates derivatives.
// Note that ystart is replaced by the corresponding values at the 
// end of the integration interval.  
     
{
  int i,k;
  int nstep;
  double x,h,dxlast;
  double *v,*vout,*dv;
  
  v   =dvector(1,nvar);
  vout=dvector(1,nvar);
  dv  =dvector(1,nvar);

  for (i=1;i<=nvar;i++){        // Load starting values
    v[i]=ystart[i];
  }
  x=x1;                         // Start integration at x=x1
  h=dx;
  nstep=(int)(floor((x2-x1)/dx+1e-8));  // Compute number of steps with dx
  for (k=1;k<=nstep;k++) {      // Take nstep steps
    (*derivs)(x,nvar,v,dv);
    rk4(v,dv,nvar,x,h,vout,derivs);
    if (h <= 1e-16) {
      cout << "Step size too small in routine RKFIXED";
    }
    x+=h;
    for (i=1;i<=nvar;i++){
      v[i]=vout[i];
    }
    normdBproc((int)(nvar/2),v,normdB);
    eraseAndDrawGraph((int)(nvar/2), nvec, normdB, x, normdBold);    
  }
 
  dxlast=x2-x1-dx*nstep;         // Take last step to x2
  if (dxlast > 1e-8) {
    (*derivs)(x,nvar,v,dv);
    rk4(v,dv,nvar,x,dxlast,vout,derivs);
    x+=dxlast;
    for (i=1;i<=nvar;i++){
      v[i]=vout[i];
    }
  }

  for (i=1;i<=nvar;i++){
     ystart[i]=v[i];
  }

  cout.setf(ios::scientific);
  cout << "  nstep= " << nstep 
       << " dxlast= " << setprecision(10) << setw(17) << dxlast
       << " dnstep= " << setprecision(10) << setw(17) << (x2-x1)/dx 
       << endl;
  cout.unsetf(ios::scientific);
  
  free_dvector(dv,1,nvar);
  free_dvector(vout,1,nvar);
  free_dvector(v,1,nvar); 
}