#include <stdio.h>
#include <math.h>
#include "RK.h"

/*
	Runge-Kutta-Fehlberg method
	
	solve an ordinary differential equation dx/dt = f(x,t)
	where x is an n-dimensional vector consisting of n variables.

	you can find how to use this function in
		rkf-solve-t0.c
*/

/*
	solve an ordinary differential equation dx/dt = f(x,t)
	from the k-th time step to the (k+1)-th time step
	
	input:
		n	dimension of vector x
		x[]	value of vector x at the k-th time step
		diff()	procedure to compute f(x,t)
		*pt	time at the k-th time step
		dt	time step (interval)
		
		procedure diff() should be in the following form:

			diff(int n, double x[], double t, double v[])
				n	dimension of vector x
				x[]	vector (n-dimensional array)
				t	time
				v[]	array where value of f(x,t) is stored
					(n-dimensional array)

	output:
		x[]	value of vector x at the (k+1)-th time step
		*pt	time at the (k+1)-th time step
*/

extern void RungeKuttaFehlbergOneStep(int n, double x[], void diff(int, double *, double, double *), double *pt, double dt)
{
	int	i;
	double               t2,          t3,          t4,          t5,          t6;
	double	             x2[MaxVars], x3[MaxVars], x4[MaxVars], x5[MaxVars], x6[MaxVars];
	double	d1[MaxVars], d2[MaxVars], d3[MaxVars], d4[MaxVars], d5[MaxVars], d6[MaxVars];
	
	diff(n, x, *pt, d1);		/* comp. d1 */
	
	t2 = (*pt) + (1./4.)*dt;
	for(i=0;i<n;i++) {
		x2[i] = x[i] + (dt/4.)*d1[i];
	}
	diff(n, x2, t2, d2);		/* comp. d2 */
	
	t3 = (*pt) + (3./8.)*dt;
	for(i=0;i<n;i++) {
		x3[i] = x[i] + (dt/32.)*(3*d1[i] + 9*d2[i]);
	}
	diff(n, x3, t3, d3);		/* comp. d3 */
	
	t4 = (*pt) + (12./13.)*dt;
	for(i=0;i<n;i++) {
		x4[i] = x[i] + (dt/2179.)*(1932*d1[i] - 7200*d2[i] + 7296*d3[i]);
	}
	diff(n, x4, t4, d4);		/* comp. d4 */

	t5 = (*pt) + dt;
	for(i=0;i<n;i++) {
		x5[i] = x[i] + dt*((439./216.)*d1[i] - 8*d2[i] + (3680./513.)*d3[i] - (845./4104.)*d4[i]);
	}
	diff(n, x5, t5, d5);		/* comp. d5 */

	t6 = (*pt) + (1./2.)*dt;
	for(i=0;i<n;i++) {
		x6[i] = x[i] + dt*(-(8./27.)*d1[i] + 2*d2[i] - (3544./2565.)*d3[i] + (1859./4104.)*d4[i] - (11./40.)*d5[i]);
	}
	diff(n, x6, t6, d6);		/* comp. d6 */

#if DEBUG
	printf("--- d1 ---\n");	for (i=0; i<n; i++) printf("%5.2lf ", d1[i]);	printf("\n");
	printf("--- d2 ---\n");	for (i=0; i<n; i++) printf("%5.2lf ", d2[i]);	printf("\n");
	printf("--- d3 ---\n");	for (i=0; i<n; i++) printf("%5.2lf ", d3[i]);	printf("\n");
	printf("--- d4 ---\n");	for (i=0; i<n; i++) printf("%5.2lf ", d4[i]);	printf("\n");
	printf("--- d5 ---\n");	for (i=0; i<n; i++) printf("%5.2lf ", d5[i]);	printf("\n");
	printf("--- d6 ---\n");	for (i=0; i<n; i++) printf("%5.2lf ", d6[i]);	printf("\n");
#endif
	
	/* update x[] */
	for(i=0;i<n;i++) {
		x[i] += (dt)*(
			      (16./135.)     *d1[i] + 
			      (0.)           *d2[i] + 
			      (6656./12825.) *d3[i] + 
			      (28561./56430.)*d4[i] +
			      (-9./50.)      *d5[i] +
			      (2./55.)       *d6[i]
			      );
	}
	/* update *t */
	*pt += dt;
}

/*
	solve an ordinary differential equation dx/dt = f(x,t)
	computation results are stored in a file

	input:
		fd	file discripter where results are stored
		n	dimension of vector x
		x[]	initial value of vector x
		diff()	procedure to compute f(x,t)
		&t	pointer to time
		dt	time step
		kcnt	number of iteration
		
		procedure diff() should be in the following form:

			diff(int n, double x[], double t, double v[])
				n	dimension of vector x
				x[]	vector (n-dimensional array)
				t	time
				v[]	array where value of f(x,t) is stored
					(n-dimensional array)

	output:
		the following values are printed to FILE *fd:
		time x[0] x[1] ... x[n-1]
*/

static char format[] = "%lf\t";
extern void RungeKuttaFehlbergfprint(FILE *fd, int n, double x[], double t)
{
	int	i;
	
	fprintf(fd, format, t);
	for(i=0;i<n;i++) {
		fprintf(fd, format, x[i]);
	}
	fprintf(fd, "\n");
}

extern void RungeKuttaFehlberg(FILE *fd, int n, double x[], void diff(int, double *, double, double *), double *pt, double dt, int kcnt)
{
	int	k, i;
	
	if (n > MaxVars) {
		fprintf(stderr, "%d exceeds MaxVars (%d)\n", n, MaxVars);
		fprintf(stderr, "increase MaxVars in opt.h and recompile all\n");
		exit(1);
	}
	
	for(k=1; k<=kcnt; k++) {
#if DEBUG
		printf("<<<<< count = %d >>>>>\n", k);
#endif
		RungeKuttaFehlbergOneStep(n, x, diff, pt, dt);
		RungeKuttaFehlbergfprint(fd, n, x, *pt);
	}
}

/*
	Adaptive Runge-Kutta-Fehlberg formula
	
	solve an ordinary differential equation dx/dt = f(x,t).
	time interval is updated adaptively
	
	you can find how to use this function in
		rkf-solve-adp-t0.c
*/

#define SafetyRatio	0.80
#define Allowance	1e-6

#define PowerTwo(a)	((a)*(a))
#define PowerFive(a)	((a)*(a)*(a)*(a)*(a))

extern void RungeKuttaFehlbergAdaptiveOneStep(int n, double x[], void diff(int, double *, double, double *), double *pt, double *pdt)
{
	int	i;
	double  t1,          t2,          t3,          t4,          t5,          t6;
	double	x1[MaxVars], x2[MaxVars], x3[MaxVars], x4[MaxVars], x5[MaxVars], x6[MaxVars];
	double	d1[MaxVars], d2[MaxVars], d3[MaxVars], d4[MaxVars], d5[MaxVars], d6[MaxVars];
	double	xest[MaxVars];
	double	dt, error, limit;
	
	dt = (*pdt);	/* current time step */
	
	t1 = (*pt);
	for(i=0;i<n;i++) {
		x1[i] = x[i];
	}
	diff(n, x1, t1, d1);		/* comp. d1 */
	
	t2 = (*pt) + (1./4.)*dt;
	for(i=0;i<n;i++) {
		x2[i] = x[i] + (dt/4.)*d1[i];
	}
	diff(n, x2, t2, d2);		/* comp. d2 */
	
	t3 = (*pt) + (3./8.)*dt;
	for(i=0;i<n;i++) {
		x3[i] = x[i] + (dt/32.)*(3*d1[i] + 9*d2[i]);
	}
	diff(n, x3, t3, d3);		/* comp. d3 */
	
	t4 = (*pt) + (12./13.)*dt;
	for(i=0;i<n;i++) {
		x4[i] = x[i] + (dt/2179.)*(1932*d1[i] - 7200*d2[i] + 7296*d3[i]);
	}
	diff(n, x4, t4, d4);		/* comp. d4 */

	t5 = (*pt) + dt;
	for(i=0;i<n;i++) {
		x5[i] = x[i] + dt*((439./216.)*d1[i] - 8*d2[i] + (3680./513.)*d3[i] - (845./4104.)*d4[i]);
	}
	diff(n, x5, t5, d5);		/* comp. d5 */

	t6 = (*pt) + (1./2.)*dt;
	for(i=0;i<n;i++) {
		x6[i] = x[i] + dt*(-(8./27.)*d1[i] + 2*d2[i] - (3544./2565.)*d3[i] + (1859./4104.)*d4[i] - (11./40.)*d5[i]);
	}
	diff(n, x6, t6, d6);		/* comp. d6 */
	
	/* update x[] */
	for(i=0;i<n;i++) {
		x[i] += (dt)*(
			      (16./135.)     *d1[i] + 
			      (0.)           *d2[i] + 
			      (6656./12825.) *d3[i] + 
			      (28561./56430.)*d4[i] +
			      (-9./50.)      *d5[i] +
			      (2./55.)       *d6[i]
			      );
	}
	
	/* compute xest[] */
	for(i=0;i<n;i++) {
		xest[i] = x1[i] + (dt)*(
					(25./216.)   *d1[i] + 
					(0.)         *d2[i] + 
					(1408./2565.)*d3[i] + 
					(2197./4104.)*d4[i] +
					(-1./5.)     *d5[i]
					);
	}
	
	error = 0.00;
	for(i=0;i<n;i++) {
		error += PowerTwo(xest[i]-x[i]);
	}
	error = sqrt(error);
	//printf("%lf %e\n", (*pt), error);
	
	/* update *t */
	*pt += dt;
	
	/* update *dt */
	if (error > Allowance * PowerFive(SafetyRatio)) {
		limit = SafetyRatio * dt * pow(Allowance/error,0.20);
		while (*pdt > limit) {
		  *pdt = (*pdt)/2;
		}
	} else {
		limit = SafetyRatio * dt * pow(Allowance/error,0.20);
		while ( 2*(*pdt) < limit) {
		  *pdt = (*pdt)*2;
		}
	}
}

extern void RungeKuttaFehlbergAdaptivefprint(FILE *fd, int n, double x[], double t, double dt)
{
	int	i;
	
	fprintf(fd, format, t);
	for(i=0;i<n;i++) {
		fprintf(fd, format, x[i]);
	}
	fprintf(fd, format, dt);
	fprintf(fd, "\n");
}

/*
	input:
		tend	time to terminate computation
*/

extern void RungeKuttaFehlbergAdaptive(FILE *fd, int n, double x[], void diff(int, double *, double, double *), double *pt, double *pdt, double tend, int kcnt)
{
	int	k, i;
	
	if (n > MaxVars) {
		fprintf(stderr, "%d exceeds MaxVars (%d)\n", n, MaxVars);
		fprintf(stderr, "increase MaxVars in opt.h and recompile all\n");
		exit(1);
	}
	
	for(k=1; k<=kcnt; k++) {
#if DEBUG
		printf("<<<<< count = %d   dt = %lf >>>>>\n", k, *pdt);
#endif
		RungeKuttaFehlbergAdaptiveOneStep(n, x, diff, pt, pdt);
		RungeKuttaFehlbergAdaptivefprint(fd, n, x, *pt, *pdt);
		if ((*pt) > tend) break;
	}
}