#include <lal/LALDetCharHveto.h>

/*
 * Okay, now we're playing a bit fast and loose. These are defined for use with
 * the SnglBurst's next pointer. Since we don't use the SnglBurst as a linked
 * list --- leaving that instead to the glib sequence --- the next pointer is
 * hijacked for use with "masking" trigs via the Prune function. Since several
 * rounds are usually required, when we "prune" a trig via its SNR, we don't
 * want to remove and then reload it later --- that takes a long time. Instead,
 * we "mask" it by assigning its next pointer one of the below values. Both
 * are near guaranteed to crash the program if someone attempts to use them
 * in the normal fashion, and are informative if looking at them in gdb.
 *
 * In short, don't use the SnglBurst next pointer for something in the hveto 
 * program. I hereby disclaim responsibility if by some magic, the program
 * actually allocates to 0xDEADBEEFDEADBEFF.
 */
#define TRIG_MASKED (SnglBurst*)0xDEADBEEFDEADBEEF
#define TRIG_NOT_MASKED (SnglBurst*)0x0

/*
 * Scan through a list of triggers and count the instances of each trigger type
 * and count its coincidences with the target channel. The hash tables for the
 * channel count and coincidence count must already be initialized. The
 * trig_sequence parameter must contain a full list of the triggers to be
 * scanned. The chan parameter is the 'target' channel -- usually h(t). The
 * twind parameter controls the length of the coincidence window.
 *
 * coinc_type = 0: Allow all coincidences
 * coinc_type = 1: Allow only one coincidence between target channel and others
 * (e.g. if two triggers from the same channel are found in coincidence, only
 * one is recorded).
 */
void XLALDetCharScanTrigs( LALGHashTable *chancount, LALGHashTable *chanhist, LALGSequence* trig_sequence, const char* chan, double twind, int coinctype ){

	/*
	 * pointer to the position in the list for the current auxiliary trigger
	 * being examined
	 */
	GSequenceIter *cur_trig = XLALGSequenceBeginRaw( trig_sequence, LALGTYPE_SNGL_BURST );

	LIGOTimeGPS start;
	LIGOTimeGPS stop;
	XLALGPSSetREAL8( &start, 0 );
	XLALGPSSetREAL8( &stop, 1 );

	// coincidence window
	LALSeg *wind;
	wind = XLALSegCreate( &start, &stop, 0 );

	// trigger pointer
	SnglBurst *sb_target;
	GHashTable *prevcoinc;

	// Iterate through trigger list -- it should be sorted by GPS time
	// TODO: Use foreach here instead?
	for( ; !g_sequence_iter_is_end(cur_trig) ; cur_trig = g_sequence_iter_next(cur_trig) ){

		// Current trigger
		sb_target = (SnglBurst*)g_sequence_get( cur_trig );
		if( sb_target->next == TRIG_MASKED ) continue;

		/*
		 * Increment the trigger count for this channel.
		 * If it is not in the hash table, add it and initialize properly
		 */
		int *value;
		value = XLALGetGHashTblIntPtr( chancount, sb_target->channel );
		if( value ){
			(*value)++;
			XLALPrintInfo( "Count: Incrementing, %s value: %i\n", sb_target->channel, *value );
			//g_hash_table_insert( chancount, XLALStringDuplicate(sb_target->channel), value );
		} else {
			XLALPrintInfo( "Count: Adding %s with time %d.%d\n", sb_target->channel, sb_target->peak_time.gpsSeconds, sb_target->peak_time.gpsNanoSeconds );
                        XLALSetGHashTblInt( chancount, sb_target->channel, 1 );
                        value = XLALGetGHashTblIntPtr( chancount, sb_target->channel );
		}

		// Is it the channel we're looking at?
		// TODO: Consider doing this from the perspective of h(t), rather than
		// each aux channel, since we're checking coincidences of N to 1 rather
		// than 1 to N.
		// FIXME: check to make sure we don't have channels which are identical
		// up to a point in the string
		if( strstr( sb_target->channel, chan ) ){
			// Yes, create window segment
			start = stop = sb_target->peak_time;
			start = *XLALGPSAdd( &start, -twind/2.0 );
			stop = *XLALGPSAdd( &stop, twind/2.0 );
			XLALPrintInfo( "creating segment from %s %d.%d %d.%d\n", sb_target->channel, start.gpsSeconds, start.gpsNanoSeconds, stop.gpsSeconds, stop.gpsNanoSeconds  );
			XLALSegSet( wind, &start, &stop, sb_target->event_id );
		} else { // No, go to the next
			continue;
		}

		// FIXME: Free memory
		prevcoinc = g_hash_table_new( g_str_hash, g_str_equal );

		XLALPrintInfo( "Checking for event %d within %d %d\n", sb_target->peak_time.gpsSeconds, wind->start.gpsSeconds, wind->end.gpsSeconds );

		// This is our secondary pointer which counts out from the current
		// trigger
		GSequenceIter *trigp = cur_trig;
		SnglBurst* sb_aux;
		gboolean begin = FALSE;
		if( g_sequence_iter_is_begin(trigp) ){
			begin = TRUE;
		} else {
			trigp = g_sequence_iter_prev(trigp);
			sb_aux = (SnglBurst*)g_sequence_get(trigp);
		}

		// Sweep backward, accumulate triggers in the window until we're outside
		// of it
		while( !begin && XLALGPSInSeg( &sb_aux->peak_time, wind ) == 0 ){

			/*
			 * If we want unique coincidences, check to see if this channel has
			 * already been added.
			 *
			 * For now, don't use the target channel
			 * FIXME: We may want this information in the future
			 */
			if( (coinctype == 1 && g_hash_table_lookup( prevcoinc, sb_aux->channel )) || strstr( sb_aux->channel, chan ) || sb_aux->next == TRIG_MASKED ){
				if( g_sequence_iter_is_begin(trigp) ) break;
				trigp = g_sequence_iter_prev(trigp);
				sb_aux = (SnglBurst*)g_sequence_get(trigp);
				continue;
			} else {
				// FIXME: Use g_hash_table_add when compatible
				g_hash_table_insert( prevcoinc, &sb_aux->channel, &sb_aux->channel );
			}

			// TODO: Macroize?
			value = XLALGetGHashTblIntPtr( chanhist, sb_aux->channel );
			// If we have an entry for this channel, use it, otherwise create a
			// new one
			if( value != NULL ){
				(*value)++;
				XLALPrintInfo( "Left Coincidence: Incrementing, %s->%ld, time %d.%d value: %i\n", sb_aux->channel, sb_target->event_id, sb_aux->peak_time.gpsSeconds, sb_aux->peak_time.gpsNanoSeconds, *value );
				//g_hash_table_insert( chanhist, &sb_aux->channel, value );
			} else {
				XLALPrintInfo( "Left Coincidence: Adding %s->%ld with time %d.%d\n", sb_aux->channel, sb_target->event_id, sb_aux->peak_time.gpsSeconds, sb_aux->peak_time.gpsNanoSeconds );
                                XLALSetGHashTblInt( chanhist, sb_aux->channel, 1 );
                                value = XLALGetGHashTblIntPtr( chanhist, sb_aux->channel );
			}
			if( g_sequence_iter_is_begin(trigp) ) break;
			trigp = g_sequence_iter_prev(trigp);
			sb_aux = (SnglBurst*)g_sequence_get(trigp);
		}

		trigp = cur_trig;
		trigp = g_sequence_iter_next(trigp);
		if( g_sequence_iter_is_end(trigp) ) break;
		sb_aux = (SnglBurst*)g_sequence_get(trigp);

		// Sweep forward, accumulate triggers in the window until we're outside
		// of it
		//do {
		while( XLALGPSInSeg( &sb_aux->peak_time, wind ) == 0 ){
			//trigp = g_sequence_iter_next(trigp);
			//if( g_sequence_iter_is_end(trigp) ) break;
			//sb_aux = (SnglBurst*)g_sequence_get(trigp);
			
			// Not the target channel?
			if( strstr( sb_aux->channel, chan ) || sb_aux->next == TRIG_MASKED ){ 
				trigp = g_sequence_iter_next(trigp);
				if( g_sequence_iter_is_end(trigp) ) break;
				sb_aux = (SnglBurst*)g_sequence_get(trigp);
				continue;
			}

			if( coinctype == 1 && g_hash_table_lookup( prevcoinc, sb_aux->channel ) ){
				trigp = g_sequence_iter_next(trigp);
				if( g_sequence_iter_is_end(trigp) ) break;
				sb_aux = (SnglBurst*)g_sequence_get(trigp);
				continue;
			} else {
				// FIXME: Use g_hash_table_add when compatible
				g_hash_table_insert( prevcoinc, &sb_aux->channel, &sb_aux->channel );
			}

			// TODO: Macroize?
			value = XLALGetGHashTblIntPtr( chanhist, sb_aux->channel );
			// If we have an entry for this channel, use it, otherwise create a
			// new one
			if( value != NULL ){
				(*value)++;
				XLALPrintInfo( "Right Coincidence: Incrementing, %s->%ld, time %d.%d value: %i\n", sb_aux->channel, sb_target->event_id, sb_aux->peak_time.gpsSeconds, sb_aux->peak_time.gpsNanoSeconds, *value );
				//g_hash_table_insert( chanhist, &sb_aux->channel, value );
			} else {
				XLALPrintInfo( "Right Coincidence: Adding %s->%ld with time %d.%d\n", sb_aux->channel, sb_target->event_id, sb_aux->peak_time.gpsSeconds, sb_aux->peak_time.gpsNanoSeconds );
                                XLALSetGHashTblInt( chanhist, sb_aux->channel, 1 );
                                value = XLALGetGHashTblIntPtr( chanhist, sb_aux->channel );
			}
			trigp = g_sequence_iter_next(trigp);
			if( g_sequence_iter_is_end(trigp) ) break;
			sb_aux = (SnglBurst*)g_sequence_get(trigp);
		}
		//} while( XLALGPSInSeg( &sb_aux->peak_time, wind ) == 0 );
		g_hash_table_destroy(prevcoinc);
	}
	XLALFree( wind );
}

/*
 * Do a round of vetoes. In short, check the significance of each channel and
 * remove the highest signifiance channel, unless it falls below the
 * significance threshold. The winner is returned as the first parameter. The
 * chan parameter is the target channel. The t_ratio parameter is the ratio of
 * the veto window duration to the total examined livetime.
 */
double XLALDetCharVetoRound( char** winner, LALGHashTable* chancount, LALGHashTable* chanhist, const char* chan, double t_ratio ){
	double mu, sig, max_sig=-1;
	int *k;

	// No reference channel present in sequence?
	// This is mostly just here to prevent a segfault later.
	if( !XLALGHashTableKeyExists(chancount, chan) ){
		XLALPrintInfo( "Reference channel not present in channel count.\n" );
		return -1.0;
	}

	GHashTableIter iter;
	gpointer key, val;
	// Number of triggers in target channel
	int n_h = XLALGetGHashTblInt(chancount, chan);

	XLALGHashTableBeginRaw( chanhist, LALGTYPE_INT, &iter );
	// Iterate over each channel and check the significance of the
	// accumulated trigger number
	while( g_hash_table_iter_next( &iter, &key, &val ) ){
		if( strstr( key, chan ) ){
			//XLALPrintInfo( stderr, "Skipping channel %s\n", (char *)key );
			continue;
		}
		// Number of triggers in auxillary channel
		int n_aux = XLALGetGHashTblInt(chancount, key);
		mu = n_h*t_ratio*n_aux;

		k = (int*)val;

		XLALPrintInfo( "Total coincidences for channel %s: %i\n", (char *)key, *k );
		XLALPrintInfo( "Mu for channel %s: %g\n", (char *)key, mu );
		sig = XLALDetCharHvetoSignificance( mu, *k );
		XLALPrintInfo( "Significance for this channel: %g\n", sig );
		if( sig > max_sig && !strstr(chan, (char*)key) ){
			max_sig = sig;
			*winner = XLALRealloc(*winner, (strlen((char *)key) + 1) * sizeof(char));
			strcpy( *winner, (char *)key );
		}
	}
	XLALPrintInfo( "winner: %s\n", *winner );
	return max_sig;
}

void XLALDetCharPruneTrigs( LALGSequence* trig_sequence, const LALSegList* onsource, double snr_thresh, const char* refchan ){
	// FIXME: Actually, this should prune all triggers
	if( onsource->length == 0 ){
		return;
	}
	size_t i = 0;

	GSequenceIter *trigp, *tmp;
	trigp = XLALGSequenceBeginRaw( trig_sequence, LALGTYPE_SNGL_BURST );

	LALSeg onseg = onsource->segs[0];
	while( !g_sequence_iter_is_end( trigp ) ){
		SnglBurst* sb = g_sequence_get( trigp );

		int pos = XLALGPSInSeg( &sb->peak_time, &onseg );

		/*
		 * The case for this sentinel:
		 * Last trigger was within the segment previous to this
		 * This trigger is in the subsequent segment, thus we need to advance
		 * the pointer to account for it
		 */
		if( pos > 0 && i < onsource->length-1 ){
			i++;
			onseg = onsource->segs[i];
			continue;
		/*
		 * We're beyond the extent of the onsource list
		 */
		} else if( pos > 0 && i >= onsource->length ) {
			tmp = trigp;
			trigp = g_sequence_iter_next( trigp );
			g_sequence_remove(tmp);
			continue;
		}

		if( pos != 0 ){
			tmp = trigp;
			trigp = g_sequence_iter_next( trigp );
			g_sequence_remove(tmp);
			continue;
		}

		/*
		 * This is the reference channel, and we don't check it's snr to remove
		 * it. If no refchan is given, then remove anything according to 
		 * criteria.
		 */
		gboolean isrefchan = (refchan != NULL);
		if( isrefchan ) isrefchan = (strstr( refchan, sb->channel ) != NULL);

		/*
		 * Mark the trigger as unused, but don't delete it. This is a common
		 * where the snr threshold is raised, but we don't want to remove the
		 * trigger. See warnings and disclaimers at the top of the file.
		 */
		if( sb->snr < snr_thresh && !isrefchan ){
			sb->next = TRIG_MASKED;
		} else {
			sb->next = TRIG_NOT_MASKED;
		}
		trigp = g_sequence_iter_next( trigp );
	}
}

/*
 * Remove triggers in window centered around peak times for triggers from vchan.
 * This is generally called after a veto round.
 *
 * TODO: Merge vetolist creation here
 * TODO: Can we also decrement the count / coincidences efficiently here?
 */
void XLALDetCharRemoveTrigs( LALGSequence* trig_sequence, LALGSequence** tbd, const LALSeg veto, const char* vchan, const char* refchan, double snr_thresh ){

	size_t vetoed_events = 0;
	size_t de_vetoed_events = 0;
	size_t nevents = XLALGetGSequenceLength(trig_sequence);
	XLALPrintInfo( "nevents: %zu\n", nevents );
	XLALPrintInfo( "Channel to veto: %s\n", vchan );

	// Pointer to the current position in the trigger list
	GSequenceIter* trigp = XLALGSequenceBeginRaw( trig_sequence, LALGTYPE_SNGL_BURST );
	// Pointer to the trigger under examination
	SnglBurst* sb;

	// Store the triggers to be deleted
	*tbd = XLALCreateGSequence(LALGTYPE_SNGL_BURST);
	GSequenceIter* tbdit = XLALGSequenceBeginRaw( *tbd, LALGTYPE_SNGL_BURST );

	// Loop over our list, looking to remove trigs
	while( !g_sequence_iter_is_end(trigp) ){
		sb = (SnglBurst*)g_sequence_get(trigp);
		if( !sb ){
			XLALPrintError( "Invalid pointer for top level iterator!\n" );
		}

		// Is it the channel to veto on?
		if( (sb->snr < snr_thresh) | !strstr( sb->channel, vchan ) ){
			trigp = g_sequence_iter_next(trigp);
			continue; // no, move on
		}
		LALSeg *trig_veto = XLALSegCreate( &veto.start, &veto.end, (int)sb->event_id );
		// Center the window on the trigger
		XLALGPSAddGPS( &trig_veto->start, &sb->peak_time);
		XLALGPSAddGPS( &trig_veto->end, &sb->peak_time);

		GSequenceIter *st = trigp, *end = trigp;

		gboolean begin = g_sequence_iter_is_begin(st);
		if( !begin ){
			st = g_sequence_iter_prev(st);
			sb = (SnglBurst*)g_sequence_get(st);
		}

		// Backwards
		while( !begin & (XLALGPSInSeg( &sb->peak_time, trig_veto ) == 0) ){
			GSequenceIter *tmp;
			tmp = st;
			if( g_sequence_iter_is_begin(st) ){
				break;
			}
			st = g_sequence_iter_prev(st);
			g_sequence_move(tmp, tbdit);
			tbdit = g_sequence_iter_next(tbdit);
			vetoed_events++;
			XLALPrintInfo( "Backwards, deleting %s (%d.%d) id %ld\n", sb->channel, sb->peak_time.gpsSeconds, sb->peak_time.gpsNanoSeconds, sb->event_id );
			if( strstr(refchan, sb->channel) ){
				de_vetoed_events++;
			}
			sb = (SnglBurst*)g_sequence_get(st);
		}

		// Check to make sure that we're not at the end of the list
		if( g_sequence_iter_is_end(end) ){
			break;
		} else {
			sb = (SnglBurst*)g_sequence_get(end);
		}

		GSequenceIter *tmp;
		// Forwards
		while( XLALGPSInSeg( &sb->peak_time, trig_veto ) == 0 ){
			// don't invalidate the top level iterator
			tmp = end;
			end = g_sequence_iter_next(end);
			// Delete this trigger -- but don't delete anything we'd use later
			if( !strstr(sb->channel, vchan) ){
				g_sequence_move(tmp, tbdit);
				tbdit = g_sequence_iter_next(tbdit);
				vetoed_events++;
				XLALPrintInfo( "Forwards, deleting %s (%d.%d) id %ld\n", sb->channel, sb->peak_time.gpsSeconds, sb->peak_time.gpsNanoSeconds, sb->event_id );
				if( strstr(refchan, sb->channel) ){
					de_vetoed_events++;
				}
			}

			if( g_sequence_iter_is_end(end) ){
				break;
			}
			sb = (SnglBurst*)g_sequence_get(end);
		}
		tmp = trigp;
		sb = (SnglBurst*)g_sequence_get(tmp);
		trigp = g_sequence_iter_next(trigp);
		g_sequence_move( tmp, tbdit );
		tbdit = g_sequence_iter_next(tbdit);
		vetoed_events++;
		XLALPrintInfo( "Veto trig, deleting %s (%d.%d) id %ld\n", sb->channel, sb->peak_time.gpsSeconds, sb->peak_time.gpsNanoSeconds, sb->event_id );

		// FIXME: Add to veto list
		XLALFree( trig_veto );

		XLALPrintInfo( "%zu events deleted so far.\n", vetoed_events );
		nevents = XLALGetGSequenceLength(trig_sequence);
		XLALPrintInfo( "%zu events remain\n", nevents );
	}
	XLALPrintInfo( "Done, total events removed %zu\n", vetoed_events );
	XLALPrintInfo( "Done, ref channel total events removed %zu\n", de_vetoed_events );

}

/*
 * Turn all the peak times for channel vchan into a segment list of vetoes.
 */
void XLALDetCharTrigsToVetoList( LALSegList* vetoes, LALGSequence* trig_sequence, const LALSeg veto, const char* vchan ){

	float wind = XLALGPSDiff(&veto.end, &veto.start);

	// Pointer to the current position in the trigger list
	GSequenceIter* trigp = XLALGSequenceBeginRaw( trig_sequence, LALGTYPE_SNGL_BURST );
	// Pointer to the trigger under examination
	SnglBurst* sb;

	while( !g_sequence_iter_is_end(trigp) ){
		sb = (SnglBurst*)g_sequence_get(trigp);
		if( !sb ){
			XLALPrintError( "Invalid pointer for top level iterator!\n" );
		}

		if( strstr( sb->channel, vchan ) ){
            LALSeg vetotmp;
            LIGOTimeGPS start = sb->peak_time;
            LIGOTimeGPS stop = sb->peak_time;
            XLALGPSSetREAL8( &start, -wind/2.0 );
            XLALGPSSetREAL8( &stop, wind/2.0 );
            XLALSegSet( &vetotmp, &start, &stop, sb->event_id );
			XLALSegListAppend( vetoes, &vetotmp );
		}
		trigp = g_sequence_iter_next(trigp);
	}
}

/*
 * Calculate the significance of a set of triggers from the Poisson survival
 * function given the expected number of triggers.
 */
double XLALDetCharHvetoSignificance( double mu, int k ){
	if( mu < 0 ){
		XLALPrintError( "%s(): attempt to calculate significance with a negative mu (%f)", __func__, mu );
	}
	if( k < 0 ){
		XLALPrintWarning( "%s(): attempt to calculate significance with a negative k (%d)", __func__, k );
	}

	double sig = gsl_sf_gamma_inc_P(k, mu);
	if( sig != 0.0 ){
		return -log10(sig);
	} else {
		return -k*log10(mu) + mu*log10(exp(1)) + gsl_sf_lngamma(k+1)/log(10);
	}
}
