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

/*
  Runge-Kutta 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 RKOneStep
(double *pt,  double dt, int n, double q[],
void diff(double, int, double *, double *))
{
  int	i;	
  double	q1[MaxVars], q2[MaxVars], q3[MaxVars];
  double	d1[MaxVars], d2[MaxVars], d3[MaxVars], d4[MaxVars];
  
  diff(*pt, n, q, d1);				/* comp. d1 */
  for(i=0;i<n;i++) {
    q1[i] = q[i] + (dt/2.00)*d1[i];
  }
  diff((*pt)+(dt/2.00), n, q1, d2);		/* comp. d2 */
  for(i=0;i<n;i++) {
    q2[i] = q[i] + (dt/2.00)*d2[i];
  }
  diff((*pt)+(dt/2.00), n, q2, d3);		/* comp. d3 */
  for(i=0;i<n;i++) {
    q3[i] = q[i] + dt*d3[i];
  }
  diff((*pt)+dt, n, q3, d4);			/* comp. d4 */
#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");
#endif
  /* update q[] */
  for(i=0;i<n;i++) {
    q[i] += (dt/6.00)*(d1[i] + 2*d2[i] + 2*d3[i] + d4[i]);
  }
  /* update t */
  *pt += dt;
}

static char format[] = "%lf\t";
extern void RKfprint(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
  dt	time step
  cnt	number of iteration
  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 RK(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
    RKOneStep(pt, dt, n, q, diff);
    RKfprint(fd, *pt, n, q);
  }
}