#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "rk.h"

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

/*
  solve an ordinary differential equation dq/dt = f(t,q)
  from the k-th time step to the (k+1)-th time step
  
  input:
  *pt	time at the k-th time step
  dt	time step (interval)
  n	dimension of vector q
  q[]	value of vector q at the k-th time step
  diff()	function to compute f(t,q)
  
  function diff() should be in the following form:
  
  diff(double t, int n, double q[], double dotq[])
  t	time
  n	dimension of vector q
  q[]	vector (n-dimensional array)
  dotq[]	array where value of f(t,q) is stored
		(n-dimensional array)
  
  output:
  *pt	time at the (k+1)-th time step
  q[]	value of vector q at the (k+1)-th time step
*/

extern void RKFOneStep
(double *pt, double dt, int n, double q[],
void diff(double, int, double *, double *))
{
  int	i;
  double	t2,          t3,          t4,          t5,          t6;
  double	q2[MaxVars], q3[MaxVars], q4[MaxVars], q5[MaxVars], q6[MaxVars];
  double	d1[MaxVars], d2[MaxVars], d3[MaxVars], d4[MaxVars], d5[MaxVars], d6[MaxVars];
  
  diff(*pt, n, q, d1);		/* comp. d1 */
  
  t2 = (*pt) + (1./4.)*dt;
  for(i=0;i<n;i++) {
    q2[i] = q[i] + (dt/4.)*d1[i];
  }
  diff(t2, n, q2, d2);		/* comp. d2 */
  
  t3 = (*pt) + (3./8.)*dt;
  for(i=0;i<n;i++) {
    q3[i] = q[i] + (dt/32.)*(3*d1[i] + 9*d2[i]);
  }
  diff(t3, n, q3, d3);		/* comp. d3 */
  
  t4 = (*pt) + (12./13.)*dt;
  for(i=0;i<n;i++) {
    q4[i] = q[i] + (dt/2179.)*(1932*d1[i] - 7200*d2[i] + 7296*d3[i]);
  }
  diff(t4, n, q4, d4);		/* comp. d4 */
  
  t5 = (*pt) + dt;
  for(i=0;i<n;i++) {
    q5[i] = q[i] + dt*((439./216.)*d1[i] - 8*d2[i] + (3680./513.)*d3[i] - (845./4104.)*d4[i]);
  }
  diff(t5, n, q5, d5);		/* comp. d5 */
  
  t6 = (*pt) + (1./2.)*dt;
  for(i=0;i<n;i++) {
    q6[i] = q[i] + dt*(-(8./27.)*d1[i] + 2*d2[i] - (3544./2565.)*d3[i] + (1859./4104.)*d4[i] - (11./40.)*d5[i]);
  }
  diff(t6, n, q6, 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 q[] */
  for(i=0;i<n;i++) {
    q[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;
}

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

/*
  solve an ordinary differential equation dq/dt = f(t,q)
  computation results are stored in a file
  
  input:
  fd	file discripter where results are stored
  &t	pointer to time
  cnt	number of iteration
  dt	time step
  n	dimension of vector q
  q[]	initial value of vector q
  diff()	function to compute f(t,q)
		
  function diff() should be in the following form:

  diff(double t, int n, double q[], double dotq[])
  t	time
  n	dimension of vector q
  q[]	vector (n-dimensional array)
  dotq[]	array where value of f(t,q) is stored
		(n-dimensional array)
  
  output:
  the following values are printed to FILE *fd:
  time q[0] q[1] ... q[n-1]
*/

extern void RKF(FILE *fd, double *pt, double dt, int cnt, int n, double q[],
void diff(double, int, double *, double *))
{
  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<=cnt; k++) {
#if DEBUG
    printf("<<<<< count = %d >>>>>\n", k);
#endif
    RKFOneStep(pt, dt, n, q, diff);
    RKFfprint(fd, *pt, n, q);
  }
}

/*
  adaptive Runge-Kutta-Fehlberg formula
  
  solve an ordinary differential equation dq/dt = f(t,q).
  time interval is updated adaptively
*/

#define SafetyRatio	0.80
#define Allowance	1e-6

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

extern void RKFadaptiveOneStep(double *pt, double *pdt, int n, double q[],
void diff(double, int, double *, double *))
{
  int	i;
  double t1,          t2,          t3,          t4,          t5,          t6;
  double q1[MaxVars], q2[MaxVars], q3[MaxVars], q4[MaxVars], q5[MaxVars], q6[MaxVars];
  double d1[MaxVars], d2[MaxVars], d3[MaxVars], d4[MaxVars], d5[MaxVars], d6[MaxVars];
  double	qest[MaxVars];
  double	dt, error, limit;
  
  dt = (*pdt);	/* current time step */
  
  t1 = (*pt);
  for(i=0;i<n;i++) {
    q1[i] = q[i];
  }
  diff(t1, n, q1, d1);		/* comp. d1 */
  
  t2 = (*pt) + (1./4.)*dt;
  for(i=0;i<n;i++) {
    q2[i] = q[i] + (dt/4.)*d1[i];
  }
  diff(t2, n, q2, d2);		/* comp. d2 */
  
  t3 = (*pt) + (3./8.)*dt;
  for(i=0;i<n;i++) {
    q3[i] = q[i] + (dt/32.)*(3*d1[i] + 9*d2[i]);
  }
  diff(t3, n, q3, d3);		/* comp. d3 */
  
  t4 = (*pt) + (12./13.)*dt;
  for(i=0;i<n;i++) {
    q4[i] = q[i] + (dt/2179.)*(1932*d1[i] - 7200*d2[i] + 7296*d3[i]);
  }
  diff(t4, n, q4, d4);		/* comp. d4 */
  
  t5 = (*pt) + dt;
  for(i=0;i<n;i++) {
    q5[i] = q[i] + dt*((439./216.)*d1[i] - 8*d2[i] + (3680./513.)*d3[i] - (845./4104.)*d4[i]);
  }
  diff(t5, n, q5, d5);		/* comp. d5 */
  
  t6 = (*pt) + (1./2.)*dt;
  for(i=0;i<n;i++) {
    q6[i] = q[i] + dt*(-(8./27.)*d1[i] + 2*d2[i] - (3544./2565.)*d3[i] + (1859./4104.)*d4[i] - (11./40.)*d5[i]);
  }
  diff(t6, n, q6, d6);		/* comp. d6 */
  
  /* update q[] */
  for(i=0;i<n;i++) {
    q[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 qest[] */
  for(i=0;i<n;i++) {
    qest[i] = q1[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(qest[i]-q[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 RKFadaptivefprint(FILE *fd, double t, double dt, int n, double q[])
{
  int	i;
  
  fprintf(fd, format, t);
  for(i=0;i<n;i++) {
    fprintf(fd, format, q[i]);
  }
  fprintf(fd, format, dt);
  fprintf(fd, "\n");
}

/*
  input:
  tend	time to terminate computation
*/

extern void RKFadaptive(FILE *fd,
double *pt, double *pdt, double tend, int cnt, int n, double q[],
void diff(double, int, double *, double *))
{
  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<=cnt; k++) {
#if DEBUG
    printf("<<<<< count = %d   dt = %lf >>>>>\n", k, *pdt);
#endif
    RKFadaptiveOneStep(pt, pdt, n, q, diff);
    RKFadaptivefprint(fd, *pt, *pdt, n, q);
    if ((*pt) > tend) break;
  }
}