#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "rk4e.h"

/* declare local functions */
int deriv(double, int, double *, double *, void *);
int truef(double, int, double *, double, double *);

/*------------------------------------------------------------------------------
  Main routine to test error-controlled Runge-Kutta integrator.
*/
int main(int argc, char *argv[])
{
    int ier, iy, len, ny;
    double xend, xstart;
    double *erry, *yin, *yout, *ystart, *ytrue;
    char *outfilename = 0x0;
    FILE *outfile;

    /* name of file to which to write output */
    printf("enter filename for output [CR = screen]: ");
    getline(&outfilename, &len, stdin);

    /* no filename: write to screen (stdout) */
    if (outfilename[0] == '\n') {
	outfile = stdout;
    /* read filename */
    } else {
	outfilename[strlen(outfilename) - 1] = '\0';
	/* open file */
	outfile = fopen(outfilename, "w");
	if (!outfile) {
	    fprintf(stderr, "cannot open %s for writing\n", outfilename);
	    return(1);
	}
    }

    /* length of y vector */
    ny = 3;

    /* allocate memory for various y arrays */
    ystart = malloc(ny * sizeof(double));
    ytrue = malloc(ny * sizeof(double));
    yin = malloc(ny * sizeof(double));
    yout = malloc(ny * sizeof(double));
    erry = malloc(ny * sizeof(double));
    
    /* set initial conditions */
    xstart = 0.;
    ystart[0] = .5;
    ystart[1] = 0.;
    ystart[2] = .5;

    /* set erry */
    for (iy = 0; iy < ny; iy++) erry[iy] = 1.;

    /* set final x */
    xend = 16.;

    while (1) {
	int done, iscan, istep;
	double abserr, relerr, dx, xin, xout;

	/* choose error */
	printf("enter relative and absolute error: ");
	iscan = scanf("%lg%*[, \t\n]%lg", &relerr, &abserr);
	if (iscan < 2) break;

	fprintf(outfile, "relerr = %g abserr = %g\n", relerr, abserr);

	/* initial x and y */
	xout = xstart;
	for (iy = 0; iy < ny; iy++) yout[iy] = ystart[iy];

	/* guess initial step size */
	dx = (xend - xstart) / 2.;

	/* write header to file */
	fprintf(outfile, "%6s %16s %16s", "step", "dx", "x");
	for (iy = 0; iy < ny; iy++) fprintf(outfile, " %16s", "y");
	for (iy = 0; iy < ny; iy++) fprintf(outfile, " %16s", "dy/y");
	fprintf(outfile, "\n");

	/* write initial x, y to file */
	fprintf(outfile, "%6d %16.8lg %16.8lg", 0, dx, xout);
	for (iy = 0; iy < ny; iy++) fprintf(outfile, " %16.8lg", yout[iy]);
	for (iy = 0; iy < ny; iy++) fprintf(outfile, " %16.8lg", 0.);
	fprintf(outfile, "\n");

	/* initialize number of steps taken */
	istep = 0;

	done = 0;

	/* call integrator repeatedly until done */
	while (!done) {
	    /* increment number of steps taken */
	    istep++;

	    /* update x and y for next step */
	    xin = xout;
	    for (iy = 0; iy < ny; iy++) yin[iy] = yout[iy];

	    /* ensure stepsize does not extend beyond xend */
	    if (xin + 2. * dx > xend) dx = (xend - xin) / 2.;

	    /* error-controlled Runge-Kutta step */
	    ier = rk4e(&deriv, xin, &xout, ny, yin, yout, &dx, relerr, abserr, erry, 0x0);

	    /* analytic solution */
	    truef(xout, ny, ytrue, xstart, ystart);

	    if (xout >= xend) done = 1;

	    /* every so often ... */
	    if (done || istep % 16 == 0) {
		/* ... write x, y to file */
		fprintf(outfile, "%6d %16.8lg %16.8lg", istep, dx, xout);
		for (iy = 0; iy < ny; iy++) fprintf(outfile, " %16.8lg", yout[iy]);
		for (iy = 0; iy < ny; iy++) fprintf(outfile, " %16.8lg", yout[iy] / ytrue[iy] - 1.);
		fprintf(outfile, "\n");
	    }

	    if (ier) {
		fprintf(stderr, "abort: rk4e returned error\n");
		done = 1;
	    }

	    fflush(outfile);
	}

    }

    return(ier);
}

/*------------------------------------------------------------------------------
  Test derivative routine,
  for hydrogen fast shock problem.
*/
int deriv(double x, int ny, double *y, double *dydx, void *pass)
{
    dydx[0] = - 3. * y[0];
    dydx[1] = 2. * y[0] - y[1];
    dydx[2] = y[0] + y[1];
}

/*------------------------------------------------------------------------------
  True analytic solution
  for hydrogen fast shock problem.

  Note use of c library routine expm1(x) == exp(x) - 1 to maintain precision.
  If expm1(x) is not available, an alternative is the hyperbolic sinh(x),
  related to expm1(x) by exp(x) - 1 = 2 exp(x/2) sinh(x/2).
*/
int truef(double x, int ny, double *y, double xstart, double *ystart)
{
    double dx;

    dx = x - xstart;

    y[0] = ystart[0] * exp(- 3. * dx);
    y[1] = (ystart[0] * (- expm1(-2. * dx)) + ystart[1]) * exp(- dx);
    y[2] = (ystart[0] + ystart[1]) * (- expm1(- dx)) + ystart[2];
}