/*************************************************************************
* Maximum Likelihood routines for the Bailey-Makeham model               *
*                                                                        *
*   Commissioned by: R. Clifton Bailey, HCFA                             *
*                    Henry Krakauer, HCFA HSQB                           *
*   Based on code written by                                             *
*             James P. Summe, Stop 2, NMRI, NNMC.  Bethesda, MD.         *
*   Author, Bill Rogers: Rand, Computing Resource Center                 *
*************************************************************************/

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

int idcons = -1;

int iter;                       /* Iteration count */
int maxiter=100;                /* Maximum number of iterations */
int parms;                      /* Number of parameters */
Int2 actparm=0;                 /* Actual number of parameters */
Int2 *actprm;                   /* Index numbers of actual parameters */
int reasoniter=1;               /* Indicator of convergence */
DOUBLE *means;                  /* Means of the data */
DOUBLE *variances;              /* Variances of the data */
DOUBLE *ntheta;                 /* Test value of theta (local to calc_ll) */
DOUBLE *adiag;                  /* Diagonal values of a (used in output) */
DOUBLE *q;                      /* Values of the first partials */
DOUBLE *qerror;                 /* Error Bounds for Q */
DOUBLE *a;                      /* Triangular Second Partial Matrix */
DOUBLE *acp;                    /* First derivative cross-products diagonal */
DOUBLE *fa;                     /* Factorization of A */
DOUBLE *b;                      /* Temporary matrix for diag of 2nd partials*/
DOUBLE *v;                      /* The new solution */
float *pch;                     /* Parameter changes */
float *lpch;                    /* Last parameter changes */
DOUBLE nobs;                    /* number of observations */
DOUBLE ll0;                     /* Likelihood of the reference model */
DOUBLE llr0;                    /* Log likelihood at beginning of step */
DOUBLE llr1;                    /* Log likelihood at test Marquardt step */
DOUBLE llrlast;                 /* Log likelihood on the previous step */
DOUBLE rho;                     /* Marquardt step size */
DOUBLE rho1;                    /* Trial Marquardt step size */
DOUBLE rhomax=5.0;              /* Maximum step size */
DOUBLE llropt;                  /* Log likelihood--Modified Marq optimum step */
DOUBLE ssexp=-5;                /* Sum of squares of exponential ? */
DOUBLE nsf=4;                   /* Number of signficant figures */
DOUBLE tolerance=0.0;
DOUBLE lambda=1.0;
DOUBLE lastt;
DOUBLE absch;
DOUBLE magnitude_largest_change;
DOUBLE sum_variances;
DOUBLE sp0[NSP];
int found_better_answer;
int debug=0;
int singularity;
int id_largest_change;
Int2 error_flag = 0;
int in_constant_pass=0;
int stopflag=0;
DOUBLE llch=0.0;

int firstpass=1;

/*************************************************************************
* msetup() allocate data areas for the maximization part of routine      *
*************************************************************************/

void msetup() {
	int i, j; char *spn;

	parms = NSP*vars;
	for (i=0; i<parms; ++i) {
		if (!fixed[i] || fixed[i]==3) actparm++;
	}

	/* allocate areas */
	actprm = (Int2 *) malloc(actparm*sizeof(int));
	means = (DOUBLE *) malloc(vars*sizeof(DOUBLE));
	variances = (DOUBLE *) malloc(vars*sizeof(DOUBLE));
	pch = (float *) malloc(parms*sizeof(float));
	lpch = (float *) malloc(parms*sizeof(float));
	ntheta = (DOUBLE *) malloc(parms*(sizeof(DOUBLE)));
	adiag = (DOUBLE *) malloc(actparm*(sizeof(DOUBLE)));
	q = (DOUBLE *) malloc(parms*(sizeof(DOUBLE)));
	acp = (DOUBLE *) malloc(parms*(sizeof(DOUBLE)));
	qerror = (DOUBLE *) malloc(parms*(sizeof(DOUBLE)));
	a = (DOUBLE *) malloc((parms*(parms+1L))/2L*(sizeof(DOUBLE)));
	fa = (DOUBLE *) malloc((parms*(parms+1L))/2L*(sizeof(DOUBLE)));
	b = (DOUBLE *) malloc(parms*(sizeof(DOUBLE)));
	v = (DOUBLE *) malloc(parms*(sizeof(DOUBLE)));
	if (debug || v==NULL || fa==NULL) printf("Allocated areas: %p %p %p %p %p %p %p\n",
			ntheta, q, qerror, a, fa, b, v);
	if (v==NULL || fa==NULL) {
		mxerror(5,"");
	}

	for (j=i=0; i<parms; ++i) {
		if (onestep && fixed[i]==3) fixed[i]=0;
		if (!fixed[i]) actprm[j++] = i;
	}
}

/* if we had confidence in the fixing algorithm, we could apply it earlier
   and avoid this */

int fix(j) int j; {
	int i, jj;
	fixed[actprm[j]] = 2;
	jj = j;
	while (++j<actparm) {
		actprm[j-1] = actprm[j];
		q[j-1] = q[j];
		qerror[j-1] = qerror[j];
		for (i=0; i<jj; ++i) a[((j-1)*j/2)+i]=a[(j*(j+1))/2+i];
		for (i=jj; i<j; ++i) a[((j-1)*j/2)+i]=a[(j*(j+1))/2+i+1];
	}
	--actparm;
	actprm[actparm] = 0x7FFF;
}

int mxlkhd() {

/**************** Author James P. Summe, Stop 2, NMRI, NNMC.  Bethesda, MD.

 Revised by: Bill Rogers.

 MXLKHD computes the maximum likelihood estimates of the parameters
 of the procedures PARTIAL and LL.  The estimates are obtained
 iteratively using the modified Marquadt method described in
 chapter 5 of Bard's Book, "Nonlinear Parameter Estimation",
 Academic Press, 1974.  Most variable names are taken from Bard's
 notation.

 */

	int i, j, k;
	DOUBLE *ap;
	DOUBLE t, dt, w;
	DOUBLE dv[NDV], sp[NSP];

	if (*logfile) logf = fopen(logfile,"w"); else logf=stdout;
	if (!logf) mxerror(21, logfile);

	if (!onestep) {
		while(!(i=sv_rpt(sp))) {
			while(!dsread()) {
				/* Calculate derived variables */
				for (i=0; i<ndv; ++i) dv[i] = d[dvlist[i]];
				sv_obs(dv);
			}
		}
		if (i==(-1)) mxerror(14,"");
		for (i=0; i<NSP; ++i) {
			j = vars*i+idcons;
			if (!fixed[j]) theta[j] = sp[i];
			if (fixed[j]==3) fixed[j] = 0;
			for (k=0; k<vars; ++k) if (k!=idcons && fixed[vars*i+k]==0) fixed[vars*i+k]=3;
		}
		for (actparm=i=0; i<parms; ++i) {
			if (!fixed[i]) actprm[actparm++] = i;
		}
		in_constant_pass = 1;
	}

	for (iter=0; reasoniter>=0 && !stopflag; ++iter) {

		if (logf!=stdout) {
			fclose(logf);
			logf=fopen(logfile,"a");
		}

		llr0 = ll (q, qerror, a, acp, theta, 2);
		if (iter==0 && onestep) ll0 = llr0; /* Save the first one */

		iter_report();
		
		if (yydebug) {
			for (i=0; i<vars;++i) printf("%8s %22.8e %22.8e\n", 
				ident[parm_number[i]], means[i], variances[i]);
			getchar();
		}
		
		if (debug) {
			printf("Iteration %d\n", iter);
			vprt(q,actparm,"First Derivatives");
			vprt(qerror,actparm,"Error Bounds");
			sprt(a,actparm,"Second Partials");
			vprt(acp,actparm,"Cross Products");
		}

/**********************************************************************
* The algorithm is further modified at this point to use the cross    *
* products as an alternative for the Marquardt ridge value            *
**********************************************************************/

		/* calcualte -A + ridge */
		for (ap=a-1,i=0; i<actparm; i++) {
			for (j=0; j<=i; j++) { ++ap; *ap = - *ap; }
			b[i] = ((*ap)<=0.0) ? acp[i] : *ap ;
			(*ap) += b[i]*lambda; /* add ridge (b) to diagonal */
		}

		if (debug && actparm<parms) {
			vprt(q,actparm,"First Derivatives");
			vprt(qerror,actparm,"Error Bounds");
			sprt(a,actparm,"Second Partials");
		}

		solution(); /* Note: sets v */

		if (debug) {
			vprt(v,actparm,"Solution for v");
		}

		/* Convergence Criteria:
		   2*v[j] < absolute or relative small
		   or derivative small
		 */

		for (reasoniter=0,j=0; j<actparm; ++j) {
			if (fabs(q[j]) < qerror[j]) continue;
			if (2.0*fabs(v[j]) >= pow(10.0,ssexp)) reasoniter=2;
			else if (2.0*fabs(v[j]) >= fabs(theta[actprm[j]])*pow(10.0,ssexp)) reasoniter=3;
		}
		
		if (iter>maxiter) {
			spout21("\n Terminating for maximum iterations\n"); 
			reasoniter=0;
		}
		
		if (llch < tolerance && llch > 0.0 && iter>=3 ) reasoniter = -1;	

		if (reasoniter<=0 && !in_constant_pass) {
			spout21("\n(convergence achieved)\n");
			output(a,fa);
			return 0;
		}

/*************************************************************************
* if we are done with the constant portion of the estimation, we must    *
* reset the "fixed" indicators and adjust the constant terms to reflect  *
* contributions of the starting values.                                  *
*       e.g.  y = b0 + b1 x  <==> y = (b0 - b1*C) + b1(x+C)              *
*                                                                        *
* Note: assumes that in_constant_pass==1 ==> idcons meaningful           *
*************************************************************************/

		if ((reasoniter<=0 || llch<1e-4) && in_constant_pass && iter>3) {
			ll0 = llr0;
			for (i=0; i<NSP; ++i) {
				sp0[i] = theta[i*vars+idcons];  /* save constant solution */
				t = sp0[i];
				for (j=0; j<vars; ++j) {
					if (fixed[i*vars+j]==3) fixed[i*vars+j]=0;
					if (j!=idcons) t -= theta[i*vars+j]*means[j];
				}
				if (!fixed[i*vars+idcons]) theta[i*vars+idcons] = t;
			}
			for (actparm=i=0; i<parms; ++i) {
				if (!fixed[i]) actprm[actparm++] = i;
			}
			in_constant_pass = 0;
			reasoniter = 1;
			iter = -1;
			continue;
		}

		for (found_better_answer=0; !found_better_answer; ) {

			rho1 = 1.0;
			for (llr1 = -INFINITY; llr1 < llr0;) {
				llr1 = calc_ll(rho1);
				if (llr1 < llr0) {
					rho1 *= 0.5;
					spout21("Inferior solution tried.\n");
				}
			}
            for (j=0; j<actparm; ++j) v[j] *= 0.999999*rho1;
		
/**********************************************************************
*    Using llr0, llr1, and the gradient vector, form a quadratic      *
*    approximation to the surface in the direction of the step.       *
*    Use this surface to interpolate or extrapolate to an optimum     *
*    stepsize.  Ref: Bard, page 111. Note: rhomax is the maximum      *
*    stepsize.                                                        *
*                                                                     *
*    It appears to me (WHR) that this is the place the algorithm goes *
*    wrong because it may take a step that is too large.  So lets     *
*    review the calculation.                                          *
*                                                                     *
*        q = globally known first derivative, and                     *
*        v = vector solution to A x = q                               *
*        rho = q * inv(A) * q =                                       *
*             (Assuming A=ll'')2 * expected LL change                 *
*                                                                     *
*    Of course, this is too small because of the Marquardt ridging.   *
*                                                                     *
*        rho = ll change expected if funqtion were ascending          *
*                   at a constant slope over the step v               *
*                                                                     *
*    So suppose f were quadratic, and we didn't know how much.        *
*    Then we would want rho = rho/(2*(rho+llr0-llr1))                 *
*                                                                     *
**********************************************************************/

			/* Calculate the slope and direction of the step */
			for (rho=0.0,j=0; j<actparm; ++j) rho += q[j]*v[j];
			

			if (debug) {
				sprintf(outbuf,"LL change:  expected=%22.7e, actual=%22.7e\n",
					rho/2.0, llr1-llr0); spout21o();
			}

			if (rho>0.0 && rho+llr0>llr1) {
				rho *= (0.5/(rho+llr0-llr1));
				if (rho>rhomax) rho=rhomax;
				sprintf(outbuf,"Optimal rho = %12g\n", rho); spout21o();
				llropt = calc_ll(rho);
			} else llropt = llr0;

			/* Find the local maximum and adjust lambda and rho */

			if (debug) {
				sprintf(outbuf,"Decision: llr0=%15g llr1=%15g llropt=%15g\n",
					llr0, llr1, llropt); spout21o();
				sprintf(outbuf,"Rho = %15g, lambda = %15g\n", rho, lambda);
				spout21o();
			}
			if (llropt>llr1 && llropt>llr0) {
				lambda *= pow(10.0,-rho);
				found_better_answer = 2;
				llch = llropt-llr0;
			}
			else if (llr1>llr0) {
				lambda /= 10.0; rho=1.0;
				found_better_answer = 1;
				llch = llr1 - llr0;
			} else {
				/* The local maximum is llr0.  Adjust a,
				   increase lambda, and obtain a new v.
				   This is not a better answer, so we
				   need to look more in the steepest
				   descent direction.
				*/
				printf("Having trouble taking a good step\n");
				if (ridge()) mxerror(1,"");
				solution(); continue;
			}
		}

		if (debug) printf("Optimal Rho = %22.7e\n", rho);

		/* Update theta */
		id_largest_change = -1;
		magnitude_largest_change = absch = 0.0;
		for (j=0; j<actparm; ++j) {
			 t = rho*v[j];
			 theta[actprm[j]] += t;
			 /* iteration diagnostics */
			 t *= t;
			 absch += t;
			 pch[j] = t;
			 i = actprm[j]%vars;
			 if (i==idcons) continue; /* Don't stop the constant */
			 w = (i==idcons) ? 1.0 : (variances[i]/sum_variances);
			 if (t*w>magnitude_largest_change) {
				id_largest_change = j;
				magnitude_largest_change = t*w;
			 }
		}

		/* Is Lambda too small? */
		if (lambda<1e-17) lambda = 1e-17;
        }
	return 0;
}

int ridge( ) {
		int i, j;
		if (lambda <= 1e16) {
				lambda *= 10.0;
				if (debug && actparm<=3) {
					printf("ridge: lambda = %15g\n",lambda);
					sprt(a,actparm,"A before ridge");
					vprt(b,actparm,"b before ridge");
				}
				for (i=0,j=1; j<=actparm; j++) {
						i += j;
						a[i-1] += b[j-1]*lambda*0.9;
				}
				if (debug) {
					sprt(a,actparm,"A after ridge");
				}
				return 0;
		}
		return 1;
}

void solution() {
	int j;
	int better;

	singularity = -1;
	while ((j=cholsky(a,actparm,fa))>=0){  /* Cholesky factorization of A */
	       if (j>=0) singularity = actprm[j];
		   if (debug)
			   printf("Cholesky return was %d\n", j);
	       if (ridge()) mxerror(2,ident[parm_number[actprm[j]%vars]]);
	       if (debug) sprt(a,actparm,"Revised matrix of partials");
	       if (debug) printf("lambda = %20g, rho = %20g\n", lambda, rho);
	}
	solve (fa, v, q, actparm);
}

/*************************************************************************
* calc_ll(rho) calculates the likelihood for the entire dataset          *
*   for a given value of rho.  rho is the amount of the stepsize         *
*   indicated in globally known v[] to add to theta[].  No derivatives   *
*   are calculated insofar as this is a trial value.                     *
*************************************************************************/

DOUBLE calc_ll(rho) DOUBLE rho; {
	DOUBLE logl;
	int j, jj, i;

	/* Form the new parameter estimates */

	for (jj=j=0; j<parms; ++j) {
		i = actprm[jj];
		if (i>j || jj==actparm) ntheta[j] = theta[j];
		else {
			ntheta[j] = theta[j]+rho*v[jj];
			++jj;
		}
	}

	/* Output the step size and the Marquardt parameter */

	if (debug) printf("rho=%12g, lambda=%12g\n", rho, lambda);

	logl=ll((DOUBLE *)NULL,(DOUBLE *)NULL,(DOUBLE *)NULL,(DOUBLE *)NULL,ntheta,0);
	if (debug) {
		vprt(ntheta,parms,"Trial solution");
		sprintf(outbuf,"<calc_ll> Likelihood = %12g,  Rho=%10g\n", logl, rho); spout21o();
	}
	return logl;
}

/*************************************************************************
* ll(q, qerror, a, acp, theta, NeedDerivatives) calculates the           *
*   likelihood function for the dataset given a value of the LOCALLY     *
*   DEFINED parameter theta.  Derivatives are calculated according to    *
*   NeedDerivatives                                                      *
*                                                                        *
*   This routine has knowledge of the global data structure and is       *
*   responsible for passing this down to the observation-level lf fcn.   *
*************************************************************************/

DOUBLE ll(q, qerror, a, acp, theta, NeedDerivatives)
DOUBLE *q;                              /* Derivatives (output)            */
DOUBLE *qerror;                         /* Error bounds (output)           */
DOUBLE *a;                              /* Second partials (output)        */
DOUBLE *acp;                            /* Sum of deriv^2 diag (output)    */
DOUBLE *theta;                          /* Parameters (input)              */
int NeedDerivatives;                    /* 0 = none, 1 = first, 2 = second */
{
	int i, j, k, m, ii, jj, kk, mm, v, vv;
	DOUBLE weight;
	DOUBLE *p;
	DOUBLE sp[NSP], spd[NSP], spde[NSP], spd2[(NSP*(NSP+1))/2], dv[NDV];
	DOUBLE x, xx, LL, t, LLobs;

	nobs = 0.0;

	/* zero out accumulations */
	if (NeedDerivatives) for (p=a,i=0; i<actparm; ++i) {
		if (NeedDerivatives==2) for (j=0; j<=i; ++j) *(p++) = 0.0;
		acp[i] = q[i] = qerror[i] = 0.0;
	}
	LL = 0.0;
	if (firstpass) for (i=0; i<vars; ++i) means[i]=variances[i]=0.0;

	while (1) {

		/* Get the data */
		if (dsread()) break;

		++nobs;

		/* Calculate structural parameters */
		for (k=0; k<NSP; ++k) sp[k] = 0.0;
		for (i=0; i<vars; ++i) {
			j = parm_number[i];
			x = (i==idcons)?1.0:d[j];
			if (firstpass) {
				t = (x-means[i])/nobs;
				means[i] += t;
				variances[i] += 2.0*(nobs-1.0)*t*t;
			}
			if (!in_constant_pass || i==idcons)
				for (k=0; k<NSP; ++k) sp[k] += x*theta[i+k*vars];
		}

		/* Calculate dependent variables */
		for (i=0; i<ndv; ++i) dv[i] = d[dvlist[i]];
		if (weightvar) weight = d[weightvar]; else weight=1.0;

		if (yydebug && debug) {
			/* compute likelhood and derivs in structural parameter form */
			vprt(dv,NDV,"Dependent Variables for Likelihood");
			vprt(sp,NSP,"Structural Parameters for Likelihood");
		}

		LL += (LLobs=LL_compute(dv,sp,spd,spde,spd2,&weight,NeedDerivatives));
#if 0
		if (LLobs==0.0 || LL==0.0) {
			printf("Debug: LL = %20g\n", LLobs);
			vprt(spd,NSP,"Structural Derivatives");
			vprt(spde,NSP,"Derivative Errors");
			sprt(spd2,NSP,"Structural second derivatives");
		}
#endif
		if (LL<=-INFINITY) {
			ds_reset();
 			return LL;
		}

		/* Translate structural paramters to actual parameters */
		/* Assumes structural parameters are in order */
		if (NeedDerivatives) for (i=0; i<actparm; ++i) {
			/* v is -1 all the time  ****************** */
			j = actprm[i];
			k = j%vars; v=parm_number[k];
			x = (k==idcons)? 1.0 : d[v];
			m = j/vars;
			/* printf("<sp %d> x[%d]=d[%d]=%20g\n", m, k, v, x); */
			t = spd[m] * x;
			q[i] += t;
			acp[i] += t*t;
			qerror[i] += spde[m]; /* ??? */
			/* printf("q[%d] = %20g\n", i, q[i]); */
			if (NeedDerivatives==2) for (ii=0; ii<=i; ++ii) {
				jj = actprm[ii];
				kk = jj%vars; vv = parm_number[kk];
				xx = (kk==idcons) ? 1.0 : d[vv];
				mm = jj/vars;
				a[(i*(i+1))/2+ii] += spd2[(m*(m+1))/2+mm]*x*xx;
			}
		}
	}
	if (firstpass) {
		sum_variances=0.0;
		for (i=0; i<vars; ++i) {
			variances[i] /= (nobs-1.0);
			sum_variances += (i==idcons)?1.0:variances[i];
		}
	}
	firstpass = 0;
	return LL;
}

void sprt(DOUBLE *xx, int n, char *s) {
	register int i ;
	int j, k ;
	printf ("*** %s\n", s);
	for (i=0;i<n;i++) {
		for (j=0;j<n;j++) {
			printf("%12.5e ",xx[(i<j)?(j*(j+1))/2+i:(i*(i+1))/2+j]) ;
		}
		printf("\n");
	}
}

void vprt(DOUBLE *xx, int n, char *s) {
    register int i ;
    sprintf (outbuf,"*** %s\n", s); spout21o();
    for (i=0;i<n;i++) {
		sprintf(outbuf,"%12.5e ",xx[i]) ; spout21o();
	}
    sprintf(outbuf,"\n"); spout21o();
}