/*************************************************************************
* Observation-level likelihoods for the Bailey-Makeham model             *
*                                                                        *
*   This is where the model is defined!                                  *
*                                                                        *
*   Commissioned by: R. Clifton Bailey, HCFA                             *
*                    Henry Krakauer, HCFA HSQB                           *
*   Based on code written by Jim Summe, NIH                              *
*   Author, Bill Rogers: Rand, Computing Resource Center                 *
*************************************************************************/

#include <stdio.h>
#include <math.h>
#include "bailey.h"

/*************************************************************************
* Compute starting values                                                *
*  The function sv_rpt() is called repeatedly.                           *
*    return = 0:  Call sv_obs() with each observation.                   *
*    return = 1:  Starting values are available                          *
*    return = -1:  Logical error in input was found (e.g.                *
*                    a survival less than 0                              *
*************************************************************************/

/*************************************************************************
* Starting value computations for the 3-state Logit (Welch model)        *
*************************************************************************/

#if defined(LOGIST)
static long reject_obs=0;
static DOUBLE group_mean[NDV];
static int sv_first=1;
static DOUBLE count;
void sv_obs(dv)
DOUBLE *dv;
{
	int i;

	if (sv_first) {
		sv_first = 0;
		for (i=0; i<3; ++i) group_mean[i] = 0.0;
	}

	if (dv[0]!=3.0 && dv[0]!=1.0 && dv[0]!=2.0) {++reject_obs; return;}
	++count;
	i = dv[0]-1;
	++group_mean[i];
}

int sv_rpt(sp)
DOUBLE *sp;
{
	int i;
	if (sv_first) return 0;
	if (group_mean[0]==0.) return -1;
	if (reject_obs) return(-1);
	if (group_mean[1]==0.) sp[0]=(-20.);
	else sp[0] = log(group_mean[1]/group_mean[0]);
	if (group_mean[2]==0.) sp[1]=(-20.);
	else sp[1] = log(group_mean[2]/group_mean[0]);
	return 1;
}
#endif

/*************************************************************************
* Starting value computations for the Bailey-Makeham Model               *
*************************************************************************/

#if defined(BAILEY)
static long reject_obs=0;
static int sv_first=0;
static LDOUBL present=0.0, mean=0.0;
static LDOUBL pltm=0.0, meanltm=0.0, pgtm=0.0, meangtm=0.0;

void sv_obs(dv)
DOUBLE *dv;
{
	if (sv_first==1) {
		if (dv[1]==1.0 && dv[0]<0.0) { ++reject_obs; return; }
		if (dv[1]==1.0) {
			++ present; 
			mean += dv[0];
		}
	} else {
		if (dv[0]<mean) {
			if (dv[1]==1.0) ++pltm;
			meanltm += dv[0];
		} else {
			if (dv[1]==1.0) ++pgtm;
			meangtm += (dv[0]-mean);
			++pltm;
			meanltm += mean;
		}
	}
}

int sv_rpt(sp)
DOUBLE *sp;
{
	LDOUBL t;
	if (reject_obs) return(-1);
	if (sv_first<=1) {
		if (sv_first==1) mean /= present;
		++sv_first;
		if (debug) printf("There were %g obs present, mean=%15g\n", present, mean);

		return 0;
	}
	if (debug) printf("%g observations < mean, %g observations > mean\n", pltm, pgtm);
	meanltm /= pltm;
	meangtm /= pgtm;
	sp[2] = -log(meangtm);
	t = 1.0/meanltm - 1.0/meangtm;
	sp[1]=-log(mean);
	if (t<=0.0) sp[0]=sp[2]-5.0; else sp[0] = log(t);
	return 1;
}
#endif

/*************************************************************************
* Likelihood computations for the Bailey-Makeham Model                   *
*************************************************************************/

#if defined(BAILEY)
DOUBLE LL_compute(dv,sp,spd,spe,spd2,weight,NeedDerivatives)
DOUBLE *dv, *sp, *spd, *spe, *spd2, *weight;
int NeedDerivatives;
{
	LDOUBL alpha, gamma, delta;
	LDOUBL gt, egt, apar, gpar, rag, dpar, lgprob;
	LDOUBL sparaa, sparag, sparad, spargg, spargd, spardd;
	LDOUBL aparerr, gparerr, dparerr;
	LDOUBL aparz, gparz, dparz;
	LDOUBL aparf, gparf, dparf;
	LDOUBL ez, rezoez;
	LDOUBL egi, oez;
	LDOUBL t;
	int death, i;

	/* Interpret input */
	if (fabs(sp[0])>100.0 || fabs(sp[1])>100.0 || fabs(sp[2])>100.0) return -INFINITY;
	alpha = exp(sp[0]);
	gamma = exp(sp[1]);
	delta = exp(sp[2]);
	t = dv[0];
	death = dv[1]!=0.0;
	if (!has_interval) interval = dv[2];

	if (!NeedDerivatives) { /* (was !NeedDerivatives) Exerpt important stuff */
		gt = gamma*t;
		sps(gt,&egt,&apar,&gpar);
		rag = alpha/gamma;
		apar = -rag*apar;
		rag = rag*egt;
		dpar = -delta*t;
		lgprob = dpar+apar;
		if (death) {
			sps ((LDOUBL)gamma*interval,&egi,&aparz,&gparz);
			aparz = -rag*aparz;
			gparz = rag*gparz - gt*aparz;
			dparz = -delta*interval;
			fps (-dparz-aparz,&ez,&oez);
			if (oez<=0.0) return -INFINITY;
			lgprob += log(oez);
		}
		return (lgprob*(*weight));
	}


	/* Calculate the log likelihood and derivatives

	The model is: hazard r(t) = a exp(-gammat) + delta

	log survivor R(t) = - <integral> (a exp(-gammat) + delta) dt

			  = -(a/gamma)(1-exp(-gammat)) - deltat

	log density f(t) = log r(t) + log R(t)

	The derivatives must be calculated with respect to log(a),
	log(gamma), and log(delta)

		LL = -(a/gamma)(1-exp(-gammat)) - deltat

		log a: -(a/gamma)(1-exp(-gammat))
		log gamma: a(t)exp(-gammat) + (a/gamma)(1-exp(-gammat))
		log delta: -deltat

		[log a][log a]: [log a]
		agamma: - (1/gamma^2)(1-exp(-gammat)) - (1/gamma)(t)exp(-gammat)
		gammagamma: 2(a/gamma^2)(t)exp(-gammat) - (a/gamma^3)(1-exp(-gammat))
			   + (a/gamma)(t^2)exp(-gammat)
		deltaa: 0
		deltagamma: 0
		[log delta][log delta]: [log delta]
	*/

	gt = gamma*t;
	sps(gt,&egt,&apar,&gpar);
	rag = alpha/gamma;
	apar = -rag*apar;   /* -(a/gamma)(1-exp(-gamma*t)) */
	gpar = rag*gpar;    /* (a/gamma)(1 - (1+gamma*t)exp(-gamma*t) */
	rag = rag*egt;      /* (a/gamma)exp(-gamma*t)) */
	dpar = -delta*t;    /* -delta*t */
	lgprob = dpar+apar;
	aparerr = -apar*b_n*FUZZ;
	gparerr = gpar*b_n*FUZZ;
	dparerr = -dpar*b_n*FUZZ;
	sparaa = apar;
	sparag = gpar;
	sparad = 0.0;
	spargg = rag*gt*gt - gpar;
	spargd = 0.0;
	spardd = dpar;

	/* For uncensored observations, the above, plus:
	   log (a exp(-gamma*t) + delta).  Compute numerical deriv's. */

	/*
		For unlogged parameters, analytic derivatives:
		a: exp(-gamma*t))/r(t)
		gamma: (-t)exp(-gamma*t)/r(t)
		delta: 1/r(t)

		aa: - [a:][a:]
		a*gamma: - [a:][gamma:] - (t)exp(-gamma*t)/r(t)
		gamma*gamma: - [gamma:][gamma:] + (t^2)(exp(-gamma*t))/r(t)
			= - [gamma:][gamma:](1-r(t)/exp(-gamma*t))
		delta*a: - [delta:][a:]
		delta*gamma: - [delta:][gamma:]
		delta*delta: - [delta:][delta:]

	*/

	if (death) {
		sps ((LDOUBL)gamma*interval,&egi,&aparz,&gparz);
		aparz = -rag*aparz;
		gparz = rag*gparz - gt*aparz;
		dparz = -delta*interval;
		fps (-dparz-aparz,&ez,&oez);
		if (oez<=0.0) {
			printf("gamma = %20g, interval = %20g\n", gamma, interval);
			printf("aparz = %20g, gparz=%20g, dparz=%20g\n",
				aparz, gparz, dparz);
			return (-INFINITY);
		}
		lgprob += log(oez);
		rezoez = -ez/oez;
		aparf = rezoez*aparz;
		gparf = rezoez*gparz;
		dparf = rezoez*dparz;
		apar += aparf;
		gpar += gparf;
		dpar += dparf;
		aparerr = (b_n+1.0)*(aparerr/b_n+aparf*FUZZ);
		gparerr = (b_n+1.0)*(gparerr/b_n+gparf*FUZZ);
		dparerr = (b_n+1.0)*(dparerr/b_n+dparf*FUZZ);
		ez = aparz-aparf;
		sparad = dparf*ez;
		ez += 1.0;
		sparaa += aparf*ez;
		sparag += gparf*ez;
		ez = gparz - gparf;
		spargd = dparf*ez;
		ez = ez - gt - 1.0;
		spargg += gparf*ez + rezoez*(alpha*gamma*egt*interval
			      *egi*(t+interval)-aparz*gt);
		spardd += dparf*(dparz-dparf+1.0);
	}
	if (NeedDerivatives) {
		spd[0] = apar;
		spd[1] = gpar;
		spd[2] = dpar;
		spe[0] = aparerr;
		spe[1] = gparerr;
		spe[2] = dparerr;
		spd2[0] = sparaa;
		spd2[1] = sparag;
		spd2[2] = spargg;
		spd2[3] = sparad;
		spd2[4] = spargd;
		spd2[5] = spardd;
		if ((*weight)!=1.0) {
			lgprob *= (*weight);
			for (i=0; i<3; ++i) {spd[i] *= (*weight); spe[i] *= (*weight);}
			for (i=0; i<6; ++i) {spd2[i] *= (*weight);}
		}
	}
	if (debug>1) {
		sprintf(outbuf,"llcert %11g %2d %11g %11g %11g %11g %11g\n",
		   t, death, interval, alpha, gamma, delta, lgprob);
		spout21o();
	}
	return lgprob;
}
#endif

/*************************************************************************
* Numerical derivatives and 2nd derivatives                              *
*************************************************************************/

/* forward declarations */
LDOUBL LLik();
LDOUBL double2 = 2.0;
LDOUBL epsilon = 1e-5;
#if defined(BAILEY)
DOUBLE LL_numeric(dv,sp,spd,spe,spd2)
#endif
#if defined(LOGIST)
DOUBLE LL_compute(dv,sp,spd,spe,spd2)
#endif
DOUBLE *dv, *sp, *spd, *spe, *spd2;
{
	int i, j;
	LDOUBL ll0, llp[10], llm[10], llpp, pplm, xsav, xxsav;
	ll0 = LLik(dv,sp);
	for (i=0; i<NSP; ++i) {
		xsav = sp[i];
		sp[i] = xsav + epsilon;
		llp[i] = LLik(dv,sp);
		for (j=0; j<i; ++j) {
			xxsav = sp[j];
			sp[j] = xxsav + epsilon;
			llpp = LLik(dv,sp);
			spd2[(i*(i+1))/2+j] = (llpp-llp[i]+ll0-llp[j]);
			sp[j] = xxsav;
		}
		sp[i] = xsav - epsilon;
		llm[i] = LLik(dv,sp);
		spd[i] = (llp[i]-llm[i])/(double2*epsilon);
		spd2[(i*(i+3))/2] = (llp[i]-ll0+llm[i]-ll0)/(epsilon*epsilon);
		spe[i] = epsilon*fabs(spd2[(i*(i+3))/2]);
		for (j=0; j<i; ++j) {
			xxsav = sp[j];
			sp[j] = xxsav - epsilon;
			llpp = LLik(dv,sp);
			spd2[(i*(i+1))/2+j] += (llpp-llm[i]+ll0-llm[j]);
			spd2[(i*(i+1))/2+j] /= (2.0*epsilon*epsilon);
			sp[j] = xxsav;
		}
		sp[i] = xsav;
	}
	return ll0;
}

/*************************************************************************
* Likelihood computations for the Bailey-Makeham Model                   *
*************************************************************************/

#if defined(BAILEY)
LDOUBL LLik(dv,sp)
DOUBLE *dv;
DOUBLE *sp;
{
	LDOUBL logr, logR, logR1, alpha, gamma, delta, tdt;
	alpha = exp(sp[0]);
	gamma = exp(sp[1]);
	delta = exp(sp[2]);
	logR = - (alpha/gamma)*(1.0-exp(-gamma*dv[0])) - delta*dv[0];
	if (dv[1]==1.0) {
		/* logr = log((alpha*exp(-gamma*dv[0]) + delta)*dv[2]); */
		tdt = dv[0]+dv[2] ;
		logR1 = - (alpha/gamma)*(1.0-exp(-gamma*tdt)) - delta*tdt;
		/* return logr+logR; */
		return log(exp(logR)-exp(logR1));
	}
	return logR;
}
#endif

/*************************************************************************
* Likelihood computations for the 3-state Logistic Model                 *
*************************************************************************/

#if defined(LOGIST)
LDOUBL LLik(dv,sp)
DOUBLE *dv;
DOUBLE *sp;
{
	LDOUBL logr, beta, gamma;
	int outcome;
	/* if (debug) printf("dv=%15g beta=%15g gamma=%15g\n", dv[0], sp[0], sp[1]); */
	outcome = dv[0];
	if (sp[0]>50. || sp[0]<-50. || sp[1]>50. || sp[1]<-50.) return -1e30;
	beta = exp(sp[0]);
	gamma = exp(sp[1]);
	logr = -log(1.0+beta+gamma);
	switch(outcome) {
		case 1: return logr;
		case 2: return sp[0]+logr;
		case 3: return sp[1]+logr;
	}
}
#endif