/* integrate.c
 * Integrate using a symplectic 5-stage, 4th order Runge-Kutta formula.
 * We use a fixed timestep.  Stolen from dsweet@andamooka.org (David Sweet),
 * who got the constants from _Numerical_Hamiltonian_Problems_ by
 * Sanz-Serna & Calvo (1994).
 */

#include <math.h>
#include <assert.h>
#include <malloc.h>
#include "misc.h"
#include "rk4s.h"
#include "matvec.h"


/*
** n = number of equations
** stiff_flag is actually "Hamiltonian" flag: 1 if we're integrating a
** hamiltonian system, otherwise 0.
** The problem zero must be 0.
** eps is actually the timestep.
** We use the given timestep internally and step past the time the user
** asks for, store that position and time, and then step back to the DTOUT
** they want.  When we get called again, we start from the internally
** stored position.  That way the solution is always advanced with the
** timestep the user requested.
*/ 
RK4S *Rk4sAlloc(int n, double Time, double *y, F_EVAL f, int stiff_flag,
    double dt, double zero)
{
    RK4S *l = Calloc(1, sizeof(RK4S));
    l->N = n;
    l->T = Time;
    l->userY = y;
    l->internalY = Malloc(n * sizeof(double));
    VecCopy(n, l->internalY, y);
    l->F = f;
    switch(stiff_flag)
    {
	case 0:
	case 1:
	    l->Hamiltonian = stiff_flag;
	    break;
	default:
	    Fatal("Rk4sAlloc: Hamiltonian flag must be 0 or 1, not %d",
		stiff_flag);
	    break;
    }

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

    l->DT = dt;

    return l;
}


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


/* Take an internal step using the constants in the given RK4S,
 * but do not modify anything in it.  Instead, make the
 * modifications to the parameter y.
 */
static void Rk4sStep(double *y, RK4S *l, double dt)
{
    static double _b[5] = { 0.0617588581356263250,
	0.3389780265536433551, 0.6147913071755775662,
	-0.1405480146593733802, 0.1250198227945261338};

    static double _B[5] = { 0.0000000000000000000,
	0.2051776615422863900, 0.4030212816042146300,
	-0.1209208763389140700, 0.5127219331924131000};

    double yp[l->N], tau = 0;
    int stage, i;

    for(stage=0; stage < 5; stage++)
    {
	/* First, update the positions.  Note that _B[0] == 0, so we
	 * can skip the first half of stage 0.
	 */
	if(stage > 0)
	{
	    tau = dt * _B[stage];
	    l->F(l->N, l->T + tau, y, yp);
	    for(i=0; i<l->N/2; i++)
		y[i] += tau * yp[i];
	}

	/* Now, update the velocities... It's unfortunate that
	 * I need to call F again because of my interface.  We
	 * really should call F "half a time" above, and the
	 * other half here, because we only update half of Y above
	 * and the other half here.  Oh well.
	 */
	tau = dt * _b[stage];
	l->F(l->N, l->T + tau, y, yp);
	for(i=l->N/2; i<l->N; i++)
	    y[i] += tau * yp[i];
    }
}



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

double Rk4sIntegrate(RK4S *l, double TOUT) /* returns actual tout */
{
    assert(TOUT >= l->T);

    /* Algorithm: keep stepping with the given timestep until such time
     * as the *next* step would put us past TOUT; don't make that last
     * step.  Instead, exit the loop and take a "special" step just to
     * assign the user's Y vector.  We don't remember this step after
     * we return.  So the user can call us again with a TOUT less than
     * our next internal step, and we'll again take a temporary step.
     */
    while(l->T + l->DT <= TOUT)
    {
	Rk4sStep(l->internalY, l, l->DT);
	l->T += l->DT;
    }

    assert(l->T <= TOUT);

    VecCopy(l->N, l->userY, l->internalY);

    if(l->T < TOUT)
	Rk4sStep(l->userY, l, TOUT - l->T);

    return TOUT;
}


void Rk4sFree(RK4S *l)
{
    free(l->internalY);
    free(l);
}
