/**************************************************************************
 *                                                                        *
 * Copyright (c) 1996 Michael Richmond and Richard Treffers               *
 *                                                                        *
 *                    This software may be copied and distributed for     *
 *                    educational, research and not for profit services   *
 *                    provided that this copyright and statement are      *
 *                    included in all such copies.                        *
 *                                                                        *
 **************************************************************************/



	/* a routine that fits a gaussian to some data, without using ANY
	   Numerical Recipes routines at all.  It makes an initial guess as
	   the parameters, then iterates, trying to minimize chi-square for
	   each parameter independently in each iteration.  Seems to work
	   pretty well, even given data only on one side of the gaussian peak.

	   if it the fit doesn't converge, it returns SDSS_ERROR and the values
	   of the parameters are left untouched.  If it does converge, the
	   function returns SDSS_OK and the values are placed into the 
	   passed parameters.   MWR 12/7/92. */


	/* 12/8/92
	 *           - changed name of function from 'gauss_fit' to 'gauss_fit'
	 *           - now include "utils.h"
	 *               ... MWR
	 */

	/* 12/11/92
	 *           - modified version for non-SDSS use in XVista.
	 *               ... MWR
	 8/24/93 - check for maximum value of exp function -rrt
	 */

#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include "pcvista.h"

#define MATH


#ifndef DEBUG
#undef DEBUG		/* define for diagnostic output */
#endif

#ifndef DEBUG2
#undef DEBUG2		/* define for LOTS of output */
#endif 

#define MAXITER    20	/* max number of iterations allowed before quit */
#define TINY  1.0e-4	/* amount to which chisq should converge */
                     	/*   and also min. fractional amount by which */
                     	/*   gaussian parameters allowed to change */
#define FACTOR 0.1	/* fraction of parameters used as a step size */
                    	/*   to determine derivative empirically      */

#define MAX_ARG 10.	/* largest argument for exp function */

#ifdef PROTO
static void guess_params(float x[], float y[], int ndata, 
                       float *amp, float *center, float *width);
static float
find_chisq(float x[], float y[], float sig[], float sig2[], int sigflag, 
			int ndata, float amp, float width, float center);
#else
static void guess_params();
static float find_chisq();
#endif
	

	/* find the gaussian that best fits the given x[] and y[] data;
	   return the amplitude (A), center (m) and width (s) of the gaussian,
	   where
	                                 (x - m)^2
	                gaussian = A exp(---------)
	                                   2*s^2

	   and the formal chisq uncertainty in the fit.  Note that if the
	   user doesn't know (or care) what the sigma in each
	   y[] point is, he can pass values of 1.0 in each sig[] and
	   then ignore the chisq.  Only "real" values for sig[] produce
	   a meaningful chisq, but this function always returns it. */

	/* return with -1 if the user passes fewer than 3 points, or if
	   the parameters don't converge within the maximum number of
	   iterations.  If there is convergence, return 0 */

#ifdef PROTO
int 
gauss_fit(float x[], float y[], float sig[], int ndata, float *amp,
                    float *center, float *width, float *chisq)
#else
int 
gauss_fit(x, y, sig, ndata, amp, center, width, chisq)
int ndata;
float x[], y[], sig[], *amp, *center, *width, *chisq;
#endif
{
	int i, iter, sign[3], sigflag;
	float a[3], newa[3], alpha[3], deriv[3], *sig2;
	float oldchisq, x1, x2, chi;

	if (ndata < 3) {
		error(0, "gauss_fit: passed fewer than three data points");
		return(-1);
	}

	/* make some initial guesses to the parameters here */
	guess_params(x, y, ndata, &(a[0]), &(a[1]), &(a[2]));

	for (i = 0; i < 3; i++) {
		alpha[i] = fabs(a[i])*FACTOR;
	}

	if ((sig2 = (float *) malloc(ndata*sizeof(float))) == NULL) {
#ifdef DEBUG
		printf("gauss_fit: couldn't malloc for sig2 - will be slow");
#endif
		sigflag = 0;
	}
	else {
		sigflag = 1;
		for (i = 0; i < ndata; i++)
			sig2[i] = sig[i]*sig[i];
	}

	/* first, we find chisq given the guessed parameters */
	chi = find_chisq(x, y, sig, sig2, sigflag, ndata, a[0], a[1], a[2]);
	oldchisq = chi;

	iter = 0;
#ifdef DEBUG
	printf("iter %2d: a[0]=%-9.4e, a[1]=%-9.4e, a[2]=%-9.4e  chi=%-9.4e\n",
				iter, a[0], a[1], a[2], chi);
#endif

	do {

	/* now, find the change in chisq w.r.t. change in each of the three
	   parameters.  use alpha[i] as the change in parameter[i] used to find 
	   the derivative empirically */
		for (i = 0; i < 3; i++) {
			newa[i] = a[i];

#ifdef DEBUG2
			printf("  param %d: ", i);
#endif
			if (alpha[i] < TINY) {
#ifdef DEBUG2
			printf("  skipping\n");	
#endif
				continue;
			}
			a[i] -= alpha[i];
			x1 = find_chisq(x, y, sig, sig2, sigflag, ndata, a[0], a[1], a[2]);
#ifdef DEBUG2
			printf("%-9.4e %-9.4e  ", a[i], x1);
#endif
			a[i] += 2.0*alpha[i];
			x2 = find_chisq(x, y, sig, sig2, sigflag, ndata, a[0], a[1], a[2]);
#ifdef DEBUG2
			printf("%-9.4e %-9.4e  ", a[i], x2);
#endif
			a[i] -= alpha[i];		/* put a[i] back to where it started */
			deriv[i] = (x2 - x1)/(2.0*alpha[i]);

#ifdef DEBUG2
			printf("%-9.4e ", deriv[i]);
#endif

			sign[i] = ((deriv[i] > 0) ? 1 : -1);

			/* now, we want to make the derivative zero by making (x2 - x1)
			   equal to zero.  So, we want to move a[i] in the direction
			   OPPOSITE the deriv, by some amount - let's say, alpha[i].  
			   If the deriv doesn't change sign on either side of the
			   current a[i], then leave the amount alpha[i] alone and 
			   use a new value for a[i].  If it does change, we're close
			   to a minimum; therefore, then cut alpha[i] by a factor of 4 
			   and don't change a[i] in this iteration. */
			if (fabs(deriv[i]) > TINY) {
				if (((deriv[i] < 0) && (x2 < chi)) ||
			   	 ((deriv[i] > 0) && (x1 < chi))) {
					newa[i] = a[i] - sign[i]*alpha[i];
#ifdef DEBUG2
					printf("new %9.4e\n", newa[i]);
#endif
					alpha[i] = fabs(newa[i])*FACTOR;
				}
				else {
					newa[i] = a[i];	/* leave it as is... */
#ifdef DEBUG2
					printf("div %9.4e\n", newa[i]);
#endif
					alpha[i] *= 0.25;
				}
			}
			else {
#ifdef DEBUG2
				printf("sam %9.4e\n", a[i]);
#endif
			}
	
		}
		for (i = 0; i < 3; i++)
			a[i] = newa[i];

		chi = find_chisq(x, y, sig, sig2, sigflag, ndata, a[0], a[1], a[2]);

		iter++;
#ifdef DEBUG
		printf("iter %2d: a[0]=%-9.4e, a[1]=%-9.4e, a[2]=%-9.4e  chi=%-9.4e\n",
				iter, a[0], a[1], a[2], chi);
#endif

		if (fabs(chi - oldchisq) < TINY)
			break;
		oldchisq = chi;

	} while (iter < MAXITER);

	/* free up the "sig2" array, if we allocated it */
	if (sigflag == 1) {
		free(sig2);
	}

	/* If we didn't converge, just return with best numbers so far */
	/* Maybe we want to print out a warning message here? */
	if (iter == MAXITER) {
		a[0] = a[0];    /* to placate compiler */
	}

	*amp = a[0];
	*center = a[1];
	/* note that my "gauss" function doesn't use the "2*sigma^2", so we
	   have to divide here by sqrt(2) */
	*width = a[2]/sqrt(2.0);
	*chisq = chi;	

	return(0);
}

	/* guess the values of the best-fitting gaussian.  I hope we don't
	   have to be too close */

#ifdef PROTO
static void
guess_params(float x[], float y[], int ndata, 
                float *amp, float *center, float *width)
#else
static void 
guess_params(x, y, ndata, amp, center, width)
int ndata;
float x[], y[], *amp, *center, *width;
#endif
{
	int i;
	float max, min, sumy, sumxy, sumxyy;

	/* guess that the amplitude is (min - max) */
	min = y[0];
	max = y[0];
	for (i = 0; i < ndata; i++) {
		if (y[i] < min)
			min = y[i];
		if (y[i] > max)
			max = y[i];
	}
	*amp = max - min;

	/* guess that the center is given by the first moment */
	sumy = 0.0;
	sumxy = 0.0;
	for (i = 0; i < ndata; i++) {
		sumy += y[i];
		sumxy += x[i]*y[i];
	}
	if (sumy == 0.0)
		*center = 0.0;
	else
		*center = sumxy/sumy;

	/* guess that the width is sqrt(the second moment) */
	sumxyy = 0.0;
	for (i = 0; i < ndata; i++) {
		sumxyy += (x[i] - *center)*(x[i] - *center)*y[i];
	}
	if (sumy == 0.0)
		*width = 1.0;
	else {
		*width = sumxyy/sumy;
		if (*width < 0.0)
			*width = 1.0;
		else
			*width = sqrt((double) *width);
	}
}
	

	/* calculate the value of chisq given the data and the values of
	   the gaussian parameters. return the chisq.  If the argument
	   "sigflag" is equal to 1, then the "sig2" argument contains
	   sig[i] squared - use it, to speed up calculations.  If the 
	   "sigflag" argument is zero, though, actually square each element
	   of the "sig" array.  Slower, but necessary if we couldn't allocate
	   space for the "sig2" array */

#ifdef PROTO
static float
find_chisq(float x[], float y[], float sig[], float sig2[], int sigflag,
			int ndata, float amp, float center, float width)
#else
static float
find_chisq(x, y, sig, sig2, sigflag, ndata, amp, center, width)
int ndata, sigflag;
float x[], y[], sig[], sig2[], amp, center, width;
#endif
{
	int i;
	float yy, sum, arg, ex;

#ifdef DEBUG3
	printf("find_chisq: given a[0]=%-9.4e, a[1]=%-9.4e, a[2]=%-9.4e\n", amp,
		center, width);
#endif

	sum = 0;

	/*
	 * this is a hack, but if someone passes "width=0", we have to
	 * do _something_.
	 */
	if (fabs(width) < TINY) {
		width = TINY;
	}
	
	if (sigflag == 1) {
		for (i = 0; i < ndata; i++) {
			arg = (x[i] - center)/width;
			if(fabs(arg) < MAX_ARG)
				ex = exp(-arg*arg);
			else
				ex=0.;
			yy = amp*ex;
			sum += (yy - y[i])*(yy - y[i])/sig2[i];
#ifdef DEBUG3
			printf("   x=%-9.4e  y=%-9.4e yy=%-9.4e  sum=%9.4e\n", x[i], y[i], yy, sum);
#endif
		}
	}
	else {
		for (i = 0; i < ndata; i++) {
			arg = (x[i] - center)/width;
			if(fabs(arg) < MAX_ARG)
				ex = exp(-arg*arg);
			else
				ex=0.;
			yy = amp*ex;
			sum += (yy - y[i])*(yy - y[i])/(sig[i]*sig[i]);
#ifdef DEBUG3
			printf("   x=%-9.4e  y=%-9.4e yy=%-9.4e  sum=%9.4e\n", x[i], y[i], yy, sum);
#endif
		}
	}

	return(sum);
}