/* integrate.c
 * Integrate using the FORTRAN routine LSODE.
 */

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

/*
** n = number of equations
** stiff_flag: 0 = no, 1 = yes.
*/ 
LSODE *LsodeAlloc(int n, double time, double *y, F_EVAL f, int stiff_flag, double eps, double zero)
{
    LSODE *l = Calloc(1, sizeof(LSODE));
    l->N = n;
    l->T = time;
    l->Y = y;
    l->F = f;
    switch(stiff_flag)
    {
	case 0: /* non-stiff */
	    l->METHFLAG = 10;
	    l->LIW = 20;
	    l->LRW = 20+16*n;
	    break;
	case 1: /* stiff */
	    l->METHFLAG = 22;
	    l->LIW = 20 + n;
	    l->LRW = 22 +9*n + n*n;
	    break;
	default:
	    Fatal("LsodeAlloc: illegal stiff_flag %d", stiff_flag);
	    break;
    }

    l->IWORK = Malloc(l->LIW * sizeof(int));
    l->RWORK = Malloc(l->LRW * sizeof(double));

    l->ITOL = 1;
    l->RTOL = eps;
    l->ATOL = (zero == 0.0 ? 1e-30 : zero);

    l->FUNOUT = 1;
    l->ISTATE = 1;
    l->IOPT = 0;

    return l;
}


/* The function called by LSODE to compute derivatives.
 *
 * If something is horribly wrong, set (*NN) to zero, and
 * LSODE will notify the calling program.
 */


static F_EVAL F;
static LSODE *L; /* if it changes, need to restart */

static void FORT_F(int *NN, double *TT, double *YY, double *YP)
{
  (*F)(*NN, *TT, YY, YP);
}


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

double LsodeIntegrate(LSODE *l, double TOUT) /* returns actual tout */
{
    int direction = TOUT - l->T > 0 ? 1 : -1;
    double fRTOL = l->RTOL, fATOL = l->ATOL;

    if(l->F != F || l != L)	/* we're changing systems */
    {
	L = l;
	F = l->F;
	l->ISTATE = 1;	/* need to re-initialize with current data */
    }

    while(direction*(l->T - TOUT) < 0)
    {
	extern void lsode_();
	lsode_(FORT_F, &l->N, l->Y, &l->T, &TOUT, &l->ITOL, &fRTOL, &fATOL,
	    &l->FUNOUT, &l->ISTATE, &l->IOPT, l->RWORK, &l->LRW, l->IWORK,
	    &l->LIW, FORT_F, &l->METHFLAG);

	switch(l->ISTATE)
	{
	case 2:           /* successful integration */
	    assert(l->T == TOUT);
	    break;
	case -1:          /* too many steps */
	    Warning("Warning: LSODE: too many steps");
	    return l->T;
	    l->ISTATE = 2; /* Just continue */
	    break;
	case -2:          /* EPS too small */
	    fRTOL *= MIN(RFA(l->RWORK,14), 10);
	    fATOL *= MIN(RFA(l->RWORK,14), 10);
	    Warning("LSODE unhappy with tolerances: trying r=%g,a=%g",
		fRTOL, fATOL);
	    l->ISTATE = 3;
	    break;
	case -3:
	    Fatal("LsodeIntegrate: bad input detected");
	    break;
	case -4:
	    Warning("LsodeIntegrate: repeated error test failures");
#if 1
	    if(l->T == TOUT)
		l->ISTATE = 2;
	    else
		return l->T;
#else
	    l->ISTATE = 2; /* Just continue */
#endif
	    break;
	case -5:
	    Warning("integrate_lsode: repeated convergence failures--"
		"bad Jacobian or mf or tolerances?");
#if 1
	    if(l->T == TOUT)
		l->ISTATE = 2;
	    else
		return l->T;
#else
	    l->ISTATE = 2; /* Just continue */
#endif
	    break;
	case -6:
	    Fatal("integrate_lsode: solution vanished with ATOL==0."
		"  Internal error.");
	    break;
	default:
	    Fatal("AAAAHHGHHGHRHGHH!  LSODE returned bogus ISTATE!");
	    break;
	}
    }
    return l->T;
}


void LsodeFree(LSODE *l)
{
    free(l->IWORK);
    free(l->RWORK);
    free(l);
}
