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

# include <hstio.h>     /* defines HST I/O functions */
# include "calnic.h"	/* defines NICMOS data structures */
# include "calnica.h"	/* defines CALNICA data structures */

# define THRESH 4.0		/* default sigma threshold for rejection */

# define DEBUG 0
# define X1		 81
# define Y1		 63
# define X2		205
# define Y2		 63

static void fitsamps (short, float *, float *, short *, float *, float,
		      float *, float *, short *, float *, short, short);
static void linfit (float *, float *, float *, short *, short, float *, float *,
		    float *, float *);
static void RejSpikes (float *, float *, short *, float *, short, float, int *);
static void RejCRs (short *, float *, short, float, int *);

/* N_CRIDCALC: Identify cosmic ray hits in NICMOS images. This routine
** calls separate subroutines for working with single (ACCUM,RAMP,BRTOBJ)
** images or multiple (MULTIACCUM) images.
**
** Revision history:
** H.Bushouse	Feb. 1996	Written for Build 2 (Version 2.0)
** H.Bushouse	05-Jun-1997	New n_mcrid routine that performs a linear
**				fit to flux vs time and uses NOISFILE to
**				compute errors. (Version 2.3)
** H.Bushouse	23-Jul-1997	Changed linfit to return zero when given 0
**				data points and return the input value when
**				given 1 data point. (Version 3.0)
** H.Bushouse	28-Jul-1997	Changed NicInfo.bunit from scalar to vector
**				(Version 3.0)
** H.Bushouse	11-Sep-1997	Added use of nic.crthresh to allow user input
**				(Version 3.0)
** H.Bushouse	29-Sep-1997	Complete rewrite of n_mcrid - eliminated use of
**				asigclip,aerrclip,merrclip,mean,wtmean routines
**				to compute mean differences and reject samples.
**				Rejection and fit to samples are now both done
**				together using new fitsamps routine. No need
**				for NOISFILE anymore. (Version 3.0)
** H.Bushouse	15-Feb-1998	Modified n_mcrid and fitsamps for special
**				handling of MultiAccum zeroth reads when
**				first read is saturated, including special
**				handling of new	ZEROSIG and GROT DQ values
**				(Version 3.2)
** H.Bushouse	05-May-1998	Modified linfit to fix division by zero bug
**				when ngood=1 and x[goodpt]=0 (Version 3.2)
** H.Bushouse	30-Sep-1998	Modified n_mcrid to use TIME image data for
**				doing count<->countrate conversions, instead of
**				using nic->exptime, since individual pixels may
**				have different exposure times (Version 3.3)
** H.Bushouse	01-Oct-1998	Modified n_mcrid to flag initial samples based
**				on value of nic->samp_rej (Version 3.3)
** H.Bushouse	05-Oct-1998	Put sections of code in fitsamps that flag
**				electronic noise spikes and CR hits into their
**				own functions "RejSpikes" and "RejCRs".
**				Modified RejSpikes so that only the largest
**				outlier is rejected on each pass. Modified
**				fitsamps to reiterate the fit after rejecting
**				spikes before looking for for CR hits (Vsn 3.3)
** H.Bushouse	16-Sep-1999	Moved static function prototypes to head of
**				file (to satisfy DEC/Alpha/OpenVMS compiler)
**				(Version 3.3)
*/

int n_cridcalc (NicInfo *nic, MultiNicmosGroup *input,
		SingleNicmosGroup *crimage) {

/* Arguments:
**	nic	 i: NICMOS info structure
**	input	io: input image(s)
**	crimage	 o: CR reject output image (for MultiACCUM only)
*/

	/* Function definitions */
	int n_mcrid (NicInfo *, MultiNicmosGroup *, SingleNicmosGroup *);
	int n_scrid (NicInfo *, MultiNicmosGroup *);

	/* Call multi-image routine for MultiACCUM datasets */
	if (nic->obsmode == MULTIACCUM) {
	    if (n_mcrid (nic, input, crimage))
		return (status);

	/* Call single-image routine for all others */
	} else {
	    if (n_scrid (nic, input))
		return (status);
	}

	/* Successful return */
	return (status = 0);
}

/* N_SCRID: Identify cosmic ray hits in single NICMOS images.
** Detected hits are only flagged in the DQ array; SCI image pixel
** values are NOT changed. */

int n_scrid (NicInfo *nic, MultiNicmosGroup *input) {

/* Arguments:
**	nic	 i: NICMOS info structure
**	input	 i: input image
*/

	/* Routine not available yet */
	sprintf (MsgText, "CRIDCALC not yet implemented; will be skipped");
	n_warn (MsgText);

	nic->CRID.corr = OMIT;

	/* Successful return */
	return (status = 0);
}

/* N_MCRID: Identify cosmic ray hits in a stack of MultiACCUM samples.
** Detected hits are flagged in the DQ arrays of the individual samples,
** but the SCI images in the individual samples are NOT changed.
** A final single-group image is produced using only unflagged data
** from the individual samples. The accumlated counts vs. time are fitted
** with a first-order (linear) polynomial using standard linear regression.
** Outliers are identified by their distance from the fit and rejected. */

int n_mcrid (NicInfo *nic, MultiNicmosGroup *input, SingleNicmosGroup *crimage){

/* Arguments:
**	nic	 i: NICMOS info structure
**	input	io: input image stack
**	crimage	 o: CR reject output image (single-group)
*/

	/* Local variables */
	short i, j, k;			/* pixel and loop indexes */
	short nsamp;			/* number of samples for pixel */
	float current_sci;		/* current sci value */
	float current_err;		/* current err value */
	short current_dq;		/* current dq  value */
	float *sci;			/* list of sci  values for pixel */
	float *err;			/* list of err  values for pixel */
	short *dq;			/* list of dq   values for pixel */
	float *time;			/* list of time values for pixel */
	float out_sci;			/* output sci  value */
	float out_err;			/* output err  value */
	short out_dq;			/* output dq   value */
	short out_samp;			/* output samp value */
	float out_time;			/* output time value */

	/* Function definitions */
	int n_copyGroup (SingleNicmosGroup *, SingleNicmosGroup *);

	/* Allocate memory for local arrays */
	sci    = NULL;
	err    = NULL;
	dq     = NULL;
	time   = NULL;
	sci    = (float *) calloc(nic->ngroups, sizeof(float));
	err    = (float *) calloc(nic->ngroups, sizeof(float));
	dq     = (short *) calloc(nic->ngroups, sizeof(short));
	time   = (float *) calloc(nic->ngroups, sizeof(float));

	/* Initialize the output crimage by copying the first input
	** group to the crimage */
	if (n_copyGroup (crimage, &(input->group[0])))
	    return (status);

	/* Set CR reject threshold; use default if no user input */
	if (nic->crthresh == 0)
	    nic->crthresh = THRESH;

	/* Print info to processing log */
	if (nic->samp_rej > 0) {
	    sprintf (MsgText, " CRIDCALC using %g sigma rejection threshold",
		     nic->crthresh);
	    n_message (MsgText);
	    sprintf (MsgText,
		     "          and rejecting first %d samples for all pixels;",
		     nic->samp_rej);
	    n_message (MsgText);
	} else {
	    sprintf (MsgText, " CRIDCALC using %g sigma rejection threshold;",
		     nic->crthresh);
	    n_message (MsgText);
	}

	/* Loop over image array, computing mean countrate at each pixel */
	for (j=0; j<input->group[0].sci.data.ny; j++) {
	for (i=0; i<input->group[0].sci.data.nx; i++) {

	     /* Get the DQ value in the zeroth-read */
	     k = nic->ngroups-1;
	     current_dq = DQPix(input->group[k].dq.data,i,j);
	     if (current_dq & ZEROSIG)
		 current_dq -= ZEROSIG;
	     if (current_dq & GROT)
		 current_dq -= GROT;

	     /* Check for pixels that are saturated already in first read
	     ** but have OK data in the zeroth read */
	     if (nic->ZSIG.corr == PERFORM &&
		 DQPix(input->group[k-1].dq.data,i,j) & SATURATED &&
		 current_dq == 0) {

		 /* For these pixels, just set the output values equal
		 ** to what's in the input zeroth-read image */
		 Pix(crimage->sci.data,i,j) = Pix(input->group[k].sci.data,i,j);
		 Pix(crimage->err.data,i,j) = Pix(input->group[k].err.data,i,j);
		 DQSetPix(crimage->dq.data,i,j,
			  DQPix(input->group[k].dq.data,i,j));
		 Pix(crimage->smpl.data,i,j) = 1;
		 Pix(crimage->intg.data,i,j) = nic->sampzero;

		 continue;
	     }

	     /* Initialize the output image DQ value */
	     out_dq = 0;

	     /* Create list of samples for this pixel */
	     nsamp = 0;
	     for (k = nic->ngroups-1; k >= 0; k--) {

		  /* Retrieve SCI, ERR, and DQ values for this sample */
		  current_sci = Pix(input->group[k].sci.data,i,j);
		  current_err = Pix(input->group[k].err.data,i,j);
		  current_dq  = DQPix(input->group[k].dq.data,i,j);

		  /* Temporarily convert countrates back to counts */
		  if (nic->bunit[0] == COUNTRATE) {
		      current_sci *= Pix(input->group[k].intg.data,i,j);
		      current_err *= Pix(input->group[k].intg.data,i,j);
		  }

		  /* Propagate DQ values to output DQ */
		  out_dq = out_dq | current_dq;

		  /* Remove ZEROSIG and GROT bits from DQ, if present,
		  ** because those samples are OK to use here (Vsn 3.2) */
		  if (current_dq & ZEROSIG)
		      current_dq -= ZEROSIG;
		  if (current_dq & GROT)
		      current_dq -= GROT;

		  /* Temporarily flag first samp_rej samples to exclude
		  ** them from the fit (Version 3.3) */
		  if (nic->samp_rej > 0 && nsamp > 0 && nsamp <= nic->samp_rej)
		      current_dq = current_dq | RESERVED1;

		  /* Copy all values to the fitting arrays */
		  sci[nsamp]  = current_sci;
		  err[nsamp]  = current_err;
		  dq[nsamp]   = current_dq;
		  time[nsamp] = Pix(input->group[k].intg.data,i,j);

		  nsamp++;
	     }

	     /* Do iterative rejection and computation of slope */
	     fitsamps (nsamp, sci, err, dq, time, nic->crthresh,
		       &out_sci, &out_err, &out_samp, &out_time, i, j);

	     /* If at least one good sample was found, reset output DQ
	     ** value to zero. But if the original data contained a
	     ** ZEROSIG or GROT value, leave those values set in output DQ
	     ** (Version 3.2) */
	     if (out_samp > 0) {
		 current_dq = 0;
		 if (out_dq & ZEROSIG)
		     current_dq += ZEROSIG;
		 if (out_dq & GROT)
		     current_dq += GROT;
		 out_dq = current_dq;
	     }

	     /* Convert results back to counts if necessary */
	     if (nic->bunit[0] == COUNTS) {
		 out_sci *= out_time;
		 out_err *= out_time;
	     }

	     /* Store final values in output crimage */
	     Pix(crimage->sci.data,i,j)  = out_sci;
	     Pix(crimage->err.data,i,j)  = out_err;
	     DQSetPix(crimage->dq.data,i,j,out_dq);
	     Pix(crimage->smpl.data,i,j) = out_samp;
	     Pix(crimage->intg.data,i,j) = out_time;

	     /* Set CR_HIT DQ for all samples following a hit */
	     for (k = 0; k < nic->ngroups-1; k++) {
		  if (dq[k] & CR_HIT)
		      dq[k+1] = dq[k+1] | CR_HIT;
	     }

	     /* Update input DQ values for detected outliers */
	     for (k = nic->ngroups-1; k >= 0; k--) {

		  if (dq[nic->ngroups-1-k] & CR_HIT)
		      DQSetPix (input->group[k].dq.data,i,j,
				DQPix (input->group[k].dq.data,i,j) | CR_HIT);

		  if (dq[nic->ngroups-1-k] & BADPIX)
		      DQSetPix (input->group[k].dq.data,i,j,
				DQPix (input->group[k].dq.data,i,j) | BADPIX);
	     }

	} /* end of loop over nx */
	} /* end of loop over ny */

	/* Free memory allocated locally */
	free (sci);
	free (err);
	free (dq);
	free (time);

	/* Successful return */
	return (status = 0);
}

/* FITSAMPS: Fit accumulating counts vs. time to compute mean countrate,
** iteratively rejecting CR hits and refitting until no new samples are
** rejected. The rejection test is based on the input error value for
** each sample.
*/

static void fitsamps (short nsamp, float *sci, float *err, short *dq,
		      float *time, float thresh, float *out_sci,
		      float *out_err, short *out_samp, float *out_time,
		      short i, short j) {

	/* Local variables */
	int k;			/* loop index */
	int nrej;		/* number of rejected samples */
	float *tsci;		/* list of cleaned fluxes */
	float *terr;		/* list of cleaned errors */
	float *ttime;		/* list of cleaned times */
	float bad_sci;		/* sci  value of bad sample(s) */
	float bad_time;		/* time value of bad sample(s) */
	int   bad_samp;		/* number of bad samples */
	float a, siga;		/* linear fit zeropoint and error */
	float b, sigb;		/* linear fit slope and error */
	float *diff;		/* sample differences relative to fit */

	/* Initialize the output values */
	*out_sci  = 0;
	*out_err  = 0;
	*out_samp = 0;
	*out_time = 0;

	/* Allocate memory for local arrays */
	tsci  = NULL;
	terr  = NULL;
	ttime = NULL;
	diff  = NULL;
	tsci   = (float *) calloc(nsamp, sizeof(float));
	terr   = (float *) calloc(nsamp, sizeof(float));
	ttime  = (float *) calloc(nsamp, sizeof(float));
	diff   = (float *) calloc(nsamp, sizeof(float));

	/* Fit and test samples for this pixel until no new samples
	** are rejected */
	nrej = 1;
	while (nrej) {

		/* Copy non-rejected sample values to tmp arrays */
		bad_samp = 0;
		bad_time = 0;
		bad_sci  = 0;
		for (k = 0; k < nsamp; k++) {

		     /* Temporarily flag zeroth read so it's not used in fit */
		     if (k == 0) {
			 tsci[k] = sci[k];
			 terr[k] = err[k];
			 dq[k] = dq[k] | BADPIX;
			 ttime[k] = time[k];

		     /* Copy unflagged samples to tmp arrays; if a previous
		     ** sample was flagged, subtract its values from
		     ** accumulated counts and time, but use original errs. */
		     } else if (dq[k] == 0) {
			 if (bad_samp == 0) {
			     tsci[k]  = sci[k];
			     ttime[k] = time[k];
			 } else {
			     tsci[k]  = sci[k]  - bad_sci;
			     ttime[k] = time[k] - bad_time;
			 }
			 terr[k] = err[k];

		     /* Accumulate bad sci and time values for flagged samples,
		     ** but only if they're not flagged as BADPIX; BADPIX is
		     ** used to flag electronic noise which does not affect
		     ** later samples. */
		     } else if (!(dq[k] & BADPIX)) {
			 bad_sci  += sci[k]  - sci[k-1];
			 bad_time += time[k] - time[k-1];
			 bad_samp++;
		     }

		     if (DEBUG && ((i==X1-1 && j==Y1-1)||(i==X2-1 && j==Y2-1))){
		         printf (" time=%g, sci=%g, err=%g, dq=%d\n", ttime[k],
			           tsci[k], terr[k], dq[k]);
		         fflush (stdout);
		     }
		}

		/* Compute mean countrate using linear fit */
		linfit (ttime, tsci, terr, dq, nsamp, &a, &b, &siga, &sigb);
		nrej = 0;

		if (DEBUG && ((i==X1-1 && j==Y1-1)||(i==X2-1 && j==Y2-1))) {
		    printf (" slope=%g, sigma=%g, intercept=%g\n", b,sigb,a);
		    fflush (stdout);
		}

		/* Compute difference of each sample from the fit */
		for (k = 0; k < nsamp; k++) {
		     if (terr[k] != 0 && dq[k] == 0)
			 diff[k] = (tsci[k]-(a+b*ttime[k])) / terr[k];
		     else
			 diff[k] = 0;

		     if (DEBUG && ((i==X1-1 && j==Y1-1)||(i==X2-1 && j==Y2-1))){
			 printf (" diff[%d] = %g\n", k, diff[k]);
			 fflush (stdout);
		     }
		}

		/* Reject electronic noise spikes */
		RejSpikes (tsci, terr, dq, diff, nsamp, thresh, &nrej);

		/* If a spike was found,
		** reiterate the fit before looking for CR hits */
		if (nrej > 0)
		    continue;

		/* Reject CR hits */
		RejCRs (dq, diff, nsamp, thresh, &nrej);

	} /* iterate the fit */

	/* Set output SCI and ERR values */
	*(out_sci) = b;
	*(out_err) = sigb;

	/* Compute output SAMP and TIME values from non-rejected samples */
	for (k = 0; k < nsamp; k++) {
	     if (dq[k] == 0) {
		 (*out_samp)++;
		 (*out_time) = ttime[k];
	     }
	}

	/* Remove BADPIX flag that was temporarily set for the zeroth read */
	if (dq[0] & BADPIX)
	    dq[0] -= BADPIX;

	/* Free local memory */
	free (tsci);
	free (terr);
	free (ttime);
	free (diff);

}

/* LINFIT: Compute a linear fit of the form y = a + bx to a set of
** data points x, y with individual standard deviations sig. Returned
** are a, b and their respective probable uncertainties siga and sigb.
** Taken from Numerical Recipes (with modifications to handle masked data).
*/

static void linfit (float *x, float *y, float *sig, short *mask, short ndata,
		    float *a, float *b, float *siga, float *sigb) {

	/* Local variables */
	int k;
	int ngood;
	int goodpt;
	float ss, sx, sy, st2, t, wt, sxoss;

	/* Initialize accumulators and results */
	ss = 0; sx = 0; sy = 0; st2 = 0;
	*a = 0; *b = 0; *siga = 0; *sigb = 0;

	/* Count number of non-masked points */
	ngood = 0;
	for (k = 0; k < ndata; k++) {
	     if (mask[k] == 0) {
		 goodpt = k;
		 ngood++;
	     }
	}

	/* Check for trivial cases */
	if (ngood == 0)
	    return;
	else if (ngood == 1 && x[goodpt] != 0) {
	    (*b)    = y[goodpt] / x[goodpt];
	    (*sigb) = sig[goodpt] / x[goodpt];
	    return;
	}

	/* Accumulate sums */
	for (k = 0; k < ndata; k++) {
	     if (sig[k] != 0 && mask[k] == 0) {
		 wt = 1.0 / (sig[k]*sig[k]);
		 ss += wt;
		 sx += x[k]*wt;
		 sy += y[k]*wt;
	     }
	}

	if (ss != 0)
	    sxoss = sx / ss;
	else
	    sxoss = 0;

	for (k = 0; k < ndata; k++) {
	     if (sig[k] != 0 && mask[k] == 0) {
		 t = (x[k] - sxoss) / sig[k];
		 st2  += t*t;
		 (*b) += t*y[k]/sig[k];
	     }
	}

	/* Solve for b and sigb */
	if (st2 != 0) {
	    (*b) = (*b) / st2;
	    (*sigb) = sqrt ( 1.0 / st2);
	} else {
	    (*b) = 0;
	    (*sigb) = 0;
	}

	/* Solve for a and siga */
	if (ss != 0 && st2 != 0) {
	    (*a) = (sy - sx*(*b)) / ss;
	    (*siga) = sqrt ( (1.0+sx*sx/(ss*st2)) / ss);
	} else {
	    (*a) = 0;
	    (*siga) = 0;
	}

}

/* REJSPIKES: Find and flag electronic noise spikes. These show up as
** large positive or negative deviations for a single sample relative to
** the samples on either side of it.
*/

static void RejSpikes (float *tsci, float*terr, short *dq, float *diff,
		       short nsamp, float thresh, int *nrej) {

	/* Local variables */
	int   k;		/* loop index */
	float max_diff; 	/* max sample deviation */
	int   max_samp;		/* sample with max deviation */

	max_diff = 0.0; max_samp = 0;

	/* Loop over samples */
	for (k = 1; k < nsamp; k++) {

	     /* Regular samples */
	     if (k < nsamp-1) {
		 if (dq[k-1] == 0 && dq[k] == 0 && dq[k+1] == 0) {

		     /* Check for positive spikes */
		     if (diff[k]-diff[k-1] > thresh &&
			 diff[k]-diff[k+1] > thresh &&
			 tsci[k]-tsci[k-1] > thresh*terr[k-1] &&
			 tsci[k]-tsci[k+1] > thresh*terr[k+1]) {
			 if (tsci[k]-tsci[k-1] > max_diff) {
			     max_diff = tsci[k]-tsci[k-1];
			     max_samp = k;
			 }
			 (*nrej)++;
		     }

		     /* Check for negative spikes */
		     if (diff[k-1]-diff[k] > thresh &&
			 diff[k+1]-diff[k] > thresh &&
			 tsci[k-1]-tsci[k] > thresh*terr[k-1] &&
			 tsci[k+1]-tsci[k] > thresh*terr[k+1]) {
			 if (tsci[k+1]-tsci[k] > max_diff) {
			     max_diff = tsci[k+1]-tsci[k];
			     max_samp = k;
			 }
			 (*nrej)++;
		     }
		 }

	     /* Handle last sample in a special way */
	     } else {
		 if (dq[k-1] == 0 && dq[k] == 0) {

		     /* Check for positive spikes */
		     if (diff[k]-diff[k-1] > thresh &&
			 tsci[k]-tsci[k-1] > thresh*terr[k-1]) {
			 if (tsci[k]-tsci[k-1] > max_diff) {
			     max_diff = tsci[k]-tsci[k-1];
			     max_samp = k;
			 }
			 (*nrej)++;
		     }

		     /* Check for negative spikes */
		     if (diff[k-1]-diff[k] > thresh &&
			 tsci[k-1]-tsci[k] > thresh*terr[k-1]) {
			 if (tsci[k-1]-tsci[k] > max_diff) {
			     max_diff = tsci[k-1]-tsci[k];
			     max_samp = k;
			 }
			 (*nrej)++;
		     }
		 }
	     }
	}

	/* If a spike was found, flag the sample that has the max deviation */
	if (*nrej > 0)
	    dq[max_samp] = dq[max_samp] | BADPIX;

}

/* REJCRS: Find and flag Cosmic Ray hits. This is done by looking for
** large negative-to-positive going changes in the sample values from one
** sample to the next.
*/

static void RejCRs (short *dq, float *diff, short nsamp, float thresh,
		    int *nrej) {

	/* Local variables */
	int k;			/* loop index */
	float prev_diff;	/* last useable difference value */

	prev_diff = 0;

	/* Loop over samples */
	for (k = 1; k < nsamp; k++) {
	     if (dq[k-1] == 0)
		 prev_diff = diff[k-1];
	     if (dq[k] == 0) {
		 if (diff[k]-prev_diff > thresh) { 
		     dq[k] = dq[k] | CR_HIT;
		     prev_diff = diff[k];
		     (*nrej)++;
		 }
	     }
	}
}