/* Generic Runge Kutta ODE Solver; version 1.0 */
/* (c) copyright Ian Hickson 2000, distributed under the GNU GPL */

#include <string.h> 
#include "dynamic-2d-arrays.h"
#include "lists.h"
#include <assert.h> /* debugging */

typedef nlNumber pStep(nlNumber x, numberList* LastValues); /* pointer to a function */

nlNumber CoupledFirstOrderDifferentialEquation(int order, nlNumber x, numberList* LastValues, 
                                               pStep* FunctionReturningHighestOrderDifferentialOfIntegrand) {
  /* This defines what is referred to as fi(x, y0..yN) in the algorithm */
  if (order >= LastValues->width-1) {
    return FunctionReturningHighestOrderDifferentialOfIntegrand(x, LastValues);
  } else {
    return nlGet(LastValues, order+1);
  }
}

numberList* ShiftNumberList(numberList* Dest, numberList* Source, numberList* Offsets, nlNumber Divisor) {
  int Index;
  assert(Dest != NULL);
  assert(Source->width == Dest->width);
  assert(Dest->width == Offsets->width);
  /* Adds Offset[i]/Divisor to Source[i] for all i, stores answer in Dest[i] */
  /* This is used in the part of the code that calculates ki(1) to ki(4) */
  for (Index = 0; Index < Source->width; Index++) {
    nlSet(Dest, Index, nlGet(Source, Index)+nlGet(Offsets, Index)/Divisor);
  }
  return Dest; /* so that the Dest list can be passed straight to another function */
}

void rungekutta(nlNumber Start, nlNumber End, nlNumber StepSize,
                numberList* InitialValues, 
                pStep* FunctionReturningHighestOrderDifferentialOfIntegrand,
                pStep* CallBack) {
  numberList* LastValues; /* yi(x) in algorithm */
  numberList* CurrentValues; /* yi(x+h) in algorithm */
  numberList* k1; /* Temporary Values - k(1) to k(4) in algorithm */
  numberList* k2;
  numberList* k3;
  numberList* k4; 
  numberList* ShiftedValues; /* gives ShiftNumberList() a temporary list to do its work */
  nlNumber x;
  int order; /* i in algorithm */
  assert(FunctionReturningHighestOrderDifferentialOfIntegrand != NULL);
  /* Allocate Memory */
  LastValues = cloneArray2D(InitialValues);
  CurrentValues = nlAllocList(LastValues->width);
  k1 = nlAllocList(LastValues->width);
  k2 = nlAllocList(LastValues->width);
  k3 = nlAllocList(LastValues->width);
  k4 = nlAllocList(LastValues->width);
  ShiftedValues = nlAllocList(LastValues->width);
  /* do the next bit for each step of the way */
  for (x = Start; x < End; x+=StepSize) {
    /* for each order of runge-kutta, find k(1), k(2), k(3) and k(4) */
    for (order=0; order < CurrentValues->width; order++) {
      nlSet(k1, order, StepSize * CoupledFirstOrderDifferentialEquation(order, x, 
                                  LastValues,
                                  FunctionReturningHighestOrderDifferentialOfIntegrand));
    }
    for (order=0; order < CurrentValues->width; order++) {
      nlSet(k2, order, StepSize * CoupledFirstOrderDifferentialEquation(order, x+StepSize/2, 
                                  ShiftNumberList(ShiftedValues, LastValues, k1, 2),
                                  FunctionReturningHighestOrderDifferentialOfIntegrand));
    }
    for (order=0; order < CurrentValues->width; order++) {
      nlSet(k3, order, StepSize * CoupledFirstOrderDifferentialEquation(order, x+StepSize/2, 
                                  ShiftNumberList(ShiftedValues, LastValues, k2, 2),
                                  FunctionReturningHighestOrderDifferentialOfIntegrand));
    }
    for (order=0; order < CurrentValues->width; order++) {
      nlSet(k4, order, StepSize * CoupledFirstOrderDifferentialEquation(order, x+StepSize, 
                                  ShiftNumberList(ShiftedValues, LastValues, k3, 1),
                                  FunctionReturningHighestOrderDifferentialOfIntegrand));
    }
    /* Now use those to find y(x+h) for each order */
    for (order=0; order < CurrentValues->width; order++) {
      nlSet(CurrentValues, order,      /* yi(x+h) = */
            nlGet(LastValues, order) + /*   yi(x)   */
            nlGet(k1, order) / 6 +     /* + ki(1)/6 */
            nlGet(k2, order) / 3 +     /* + ki(2)/3 */
            nlGet(k3, order) / 3 +     /* + ki(3)/3 */
            nlGet(k4, order) / 6       /* + ki(4)/6 */
            );
    }
    /* Set LastValues to CurrentValues */
    nlFreeList(LastValues);
    LastValues = cloneArray2D(CurrentValues);
    /* Call the CallBack routine to do whatever it wants with the data for this value of x */
    CallBack(x, CurrentValues); /* nlGet(CurrentValues, 0) gives the value of the function at x */
    /* Adjust StepSize here if necessary [not implemented!] */
  }
  /* free the memory up again */
  nlFreeList(LastValues);
  nlFreeList(CurrentValues);
  nlFreeList(k1);
  nlFreeList(k2);
  nlFreeList(k3);
  nlFreeList(k4);
  nlFreeList(ShiftedValues);
}

/* end */