/* integrate.c
 * Integrate using classical 4-stage, 4th order Runge-Kutta formula.
 * We use a fixed timestep.
 */

#include <math.h>
#include <assert.h>
#include <malloc.h>
#include "misc.h"
#include "rk4.h"

/*
** n = number of equations
** stiff_flag must be 0, as must the problem zero.
** eps is actually the timestep.
*/ 
RK4 *Rk4Alloc(int n, double Time, double *y, F_EVAL f, int stiff_flag,
    double dt, double zero)
{
    RK4 *l = Calloc(1, sizeof(RK4));
    l->N = n;
    l->T = Time;
    l->Y = y;
    l->F = f;
    if(stiff_flag != 0)
    {
	Fatal("Rk4Alloc: stiff_flag must be 0, not %d", stiff_flag);
    }

    if(zero != 0)
    {
	Fatal("Rk4Alloc: problem zero must be 0, not %g", zero);
    }

    l->DT = dt;

    return l;
}


/* The function called by RK4 to compute derivatives.
 *
 * If something is horribly wrong, set (*NN) to zero.
 */


/****************************************************
 * integrate()
****************************************************/

double Rk4Integrate(RK4 *l, double TOUT) /* returns actual tout */
{
    int nsteps, step, i;
    const int N = l->N;
    double dt;

    assert(TOUT >= l->T);

    if(l->T == TOUT)
	return TOUT;

    /* "normalize" the timestep */
    nsteps = ceil((TOUT - l->T)/l->DT);
    dt = (TOUT - l->T)/nsteps;

    for(step=0; step < nsteps; step++)
    {
	double dx1[N], dx2[N], dx3[N], dx4[N], y[N], yp[N];

	/* First stage */

	l->F(l->N, l->T, l->Y, yp);
	for(i=0; i<N; i++)
	{
	    dx1[i] = dt * yp[i];
	    y[i] = l->Y[i] + dx1[i]/2;
	}

	/* Second stage */
	l->F(l->N, l->T+dt/2, y, yp);
	for(i=0; i<N; i++)
	{
	    dx2[i] = dt * yp[i];
	    y[i] = l->Y[i] + dx2[i]/2;
	}

	/* Third stage */
	l->F(l->N, l->T+dt/2, y, yp);
	for(i=0; i<N; i++)
	{
	    dx3[i] = dt * yp[i];
	    y[i] = l->Y[i] + dx3[i];
	}

	/* Fourth stage */
	l->F(l->N, l->T+dt, y, yp);
	for(i=0; i<N; i++)
	{
	    dx4[i] = dt * yp[i];
	    l->Y[i] += (dx1[i] + dx4[i])/6 + (dx2[i] + dx3[i])/3;
	}
	l->T += dt;
    }

    /*
    ** Make sure T is very close to tout.
    */
    assert(fabs(l->T-TOUT)/fabs(MAX(fabs(l->T),fabs(TOUT))) < 1e-15*nsteps);
    return l->T = TOUT;
}


void Rk4Free(RK4 *l)
{
    free(l);
}
