/* integrate.c
 *
 * Integrate using the FORTRAN routine DDRIV2 from the book _Numerical
 * Methods and Software_, by Kahaner, Moler & Nash.
 *
 * This file written by Wayne Hayes for a csc2529 project.
 */

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

void Ddriv2Size(DDRIV2 *d)
{
    int stiff_flag = d->MINT - 1;

    switch(stiff_flag)
    {
    case 0: /* non-stiff */
	d->LW = 16*d->N + 2*d->NROOT + 204;
	break;
    case 1: /* stiff */
	d->LW = d->N*d->N + 10*d->N + 2*d->NROOT + 204;
	break;
    case 2: /* dynamic */
	d->LW = d->N*d->N + 17*d->N + 2*d->NROOT + 204;
	break;
    default: Fatal("Ddriv2Alloc: illegal stiff_flag %d", stiff_flag);
    }
    if(d->W)
	d->W = (double*) Realloc(d->W, d->LW * sizeof(d->W[0]));
    else
	d->W = (double*) Malloc (      d->LW * sizeof(d->W[0]));

    d->LIW = d->N + 21;
    if(d->IW)
	d->IW = (int*) Realloc(d->IW, d->LIW * R8*sizeof(d->IW[0]));
    else
	d->IW = (int*) Malloc (       d->LIW * R8*sizeof(d->IW[0]));
}

/*
** n = number of equations
** stiff_flag: 0 = no, 1 = yes, 2 = dynamically decided by DDRIV2
*/ 
DDRIV2 *Ddriv2Alloc(int n, double time, double *y, F_EVAL f, int stiff_flag,
    double eps, double zero)
{
    DDRIV2 *d = Calloc(1, sizeof(DDRIV2));

    d->NROOT = 0;	/* change these by calling Ddriv2PrepRoot */
    d->G = NULL;

    d->N = n;
    d->MSTATE = 1;
    d->MINT = stiff_flag + 1;

    Ddriv2Size(d);

    d->T = time;
    d->Y = y;
    /* TOUT not set until integrate() is called */
    d->EPS = eps;
    d->EWT = zero;
    d->F = f;

    return d;
}

void Ddriv2PrepRoot(DDRIV2 *d, int nroot, G_EVAL g)
{
    d->NROOT = nroot;
    d->G = g;
    Ddriv2Size(d);
    d->MSTATE = 1;
}

int Ddriv2WhichRoot(DDRIV2 *d)
{
    return d->whichRoot;
}

void Ddriv2Free(DDRIV2 *d)
{
    Free(d->W);
    Free(d->IW);
    Free(d);
}


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

#if INTEGRATOR_PROGRESS
static double prevTime;
static integrator_progress = true;
#endif
static DDRIV2 *_gd;


static void FORT_F(int *NN, double *TT, double *YY, double *YP)
{
#if INTEGRATOR_PROGRESS
    if (integrator_progress && prevTime != *TT) {
	fprintf(stderr, "\015time %-10g, step %-10g O%d S%d   ",
	    *TT, *TT-prevTime, IFA(_gd->IW,2), IFA(_gd->IW,16)-1);
    prevTime = *TT;
  }
#endif
    (*(_gd->F))(*NN, *TT, YY, YP);
}

static double FORT_G(int *NN, double *TT, double *YY, int *IROOT)
{
    return (*(_gd->G))(*NN, *TT, YY, (*IROOT) - 1);
}


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

double Ddriv2Integrate(DDRIV2 *d, double TOUT) /* returns actual tout */
{
    int direction = TOUT - d->T > 0 ? 1 : -1;

    _gd = d;

    while(direction*(d->T - TOUT) < 0)
    {
	extern ddriv2_();
	ddriv2_(&d->N, &d->T, d->Y, FORT_F, &TOUT, &d->MSTATE, &d->NROOT,
	    &d->EPS, &d->EWT, &d->MINT, d->W, &d->LW, d->IW, &d->LIW, FORT_G);

	switch(ABS(d->MSTATE))
	{
	case 2:        /* successful integration */
	    assert(d->T == TOUT);
	    break;
	case 3:   /* too many steps */
	    /*Warning("DDRIV2 took 1000 steps; continuing anyway");*/
	    break;
	case 4:   /* EPS too small */
	    /*Warning("DDRIV2 unhappy with EPS: reset to %g", EPS);*/
	    break;
	case 5:   /* Root found. Index of which is in 6th element of IWORK */
	    assert(d->NROOT > 0);
	    d->whichRoot = IFA(d->IW,6)-1;
	    return d->T;
	    break;       
	case 6:           /* F set N to zero */
	    Fatal("F set N to zero");
	    break;
	case 7:           /* G set N to zero */
	    Fatal("G set N to zero");
	    break;
	default:
	    Fatal("ACK!  DDRIV2 returned bogus MSTATE!");
	    break;
	}

	#if INTEGRATOR_PROGRESS
	if (integrator_progress)
	    fprintf(stderr, "\015time %-10g, step %-10g O%d S%d   "
		"fullstep: avgO(%-6g), %s    ",
		d->T, d->T-prevTime, IFA(d->IW,2), IFA(d->IW,16)-1, RFA(d->W,3),
		(IFA(d->IW,16) == 1 ? "nonstiff" : "STIFF") );
	#endif
    }
    return d->T;
}
