#!/usr/bin/python
#
# Copyright (C) 2008  Kipp Cannon
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.


#
# =============================================================================
#
#                                   Preamble
#
# =============================================================================
#


import glob
import math
from optparse import OptionParser
import sys


from glue.lal import CacheEntry
from glue.ligolw import ligolw
from glue.ligolw import lsctables
from glue.ligolw import utils as ligolw_utils
from glue.ligolw.utils import process as ligolw_process
from glue.ligolw.utils import segments as ligolw_segments
from glue import segmentsUtils
from pylal import git_version
from pylal import ligolw_thinca
from pylal import snglcoinc
from pylal.xlal.datatypes.ligotimegps import LIGOTimeGPS


lsctables.use_in(ligolw.LIGOLWContentHandler)


__author__ = "Kipp Cannon <kipp.cannon@ligo.org>"
__version__ = "git id %s" % git_version.id
__date__ = git_version.date


#
# Use interning row builder to save memory.
#


lsctables.table.RowBuilder = lsctables.table.InterningRowBuilder


#
# Use C row classes for memory efficiency and speed.
#


lsctables.SnglInspiralTable.RowType = lsctables.SnglInspiral = ligolw_thinca.SnglInspiral


#
# =============================================================================
#
#                                 Command Line
#
# =============================================================================
#


def parse_command_line():
	parser = OptionParser(
		version = "Name: %%prog\n%s" % git_version.verbose_msg,
		usage = "%prog [options] [file ...]",
		description = "%prog implements the inspiral coincidence algorithm for use in performing trigger-based multi-instrument searches for gravitational wave events.  The LIGO Light Weight XML files listed on the command line are processed one by one in order, and over-written with the results.  If no files are named, then input is read from stdin and output written to stdout.  Gzipped files will be autodetected on input, if a file's name ends in \".gz\" it will be gzip-compressed on output."
	)
	parser.add_option("-c", "--comment", metavar = "text", help = "Set comment string in process table (default = None).")
	parser.add_option("-f", "--force", action = "store_true", help = "Process document even if it has already been processed.")
	parser.add_option("-t", "--e-thinca-parameter", metavar = "float", type = "float", help = "Set the ellipsoidal coincidence algorithm's threshold (required).")
	parser.add_option("--exact-mass", action = "store_true", help = "Use the ellipsoidal coincidence algorithm and require exact mass coincidence.")
	parser.add_option("-a", "--effective-snr", metavar = "{inspiral|gstlal}", help = "Select effective SNR algorithm (required).  \"inspiral\" reproduces the algorithm used by lalapps_thinca, \"gstlal\" is an algorithm appropriate for use with gstlal_inspiral.")
	parser.add_option("--effective-snr-factor", metavar = "float", type = "float", default = 250.0, help = "Set the effective SNR factor used with the \"inspiral\" effective SNR (default = 250.0).")
	parser.add_option("--vetoes-name", metavar = "string", default = "vetoes", help = "From the input document, exatract the segment list having this name to use as the veto segments (default = \"vetoes\").  Warning:  if no segments by this name are found in the document then vetoes will not be applied, this is not an error condition.")
	parser.add_option("--trigger-program", metavar = "name", default = "inspiral", help = "Set the name of the program that generated the event list as it appears in the process table (default = \"inspiral\").")
	parser.add_option("--coinc-end-time-segment", metavar = "seg", help = "The segment of time to retain coincident triggers from. Uses segmentUtils.from_range_strings() format \"START:END\" for an interval of the form [START,END), \"START:\" for an interval of the form [START,INF), and \":END\" for an interval of the form (-INF,END).")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	options, filenames = parser.parse_args()

	#
	# check arguments
	#

	required_options = ["e_thinca_parameter", "effective_snr"]
	missing_options = [option for option in required_options if getattr(options, option) is None]
	if missing_options:
		raise ValueError("missing required option(s) %s" % ", ".join("--%s" % option.replace("_", "-") for option in missing_options))
	if options.effective_snr not in ("inspiral", "gstlal"):
		raise ValueError("unrecognized --effective-snr %s" % options.effective_snr)

	if options.coinc_end_time_segment is not None:
		if ',' in options.coinc_end_time_segment:
			raise ValueError("--coinc-end-time-segment may only contain a single segment")
		options.coinc_end_time_segs = segmentsUtils.from_range_strings([options.coinc_end_time_segment], boundtype = LIGOTimeGPS).coalesce()
	else:
		options.coinc_end_time_segs = None

	#
	# done
	#

	return options, (filenames or [None])


#
# =============================================================================
#
#                           Add Process Information
#
# =============================================================================
#


process_program_name = "ligolw_thinca"


def append_process(xmldoc, comment = None, force = None, e_thinca_parameter = None, exact_mass = None, effective_snr_factor = None, vetoes_name = None, trigger_program = None, effective_snr = None, coinc_end_time_segment = None, verbose = None):
	process = ligolw_process.append_process(xmldoc, program = process_program_name, version = __version__, cvs_repository = u"lscsoft", cvs_entry_time = __date__, comment = comment)

	params = [
		(u"--e-thinca-parameter", u"real_8", e_thinca_parameter)
	]
	if comment is not None:
		params += [(u"--comment", u"lstring", comment)]
	if force is not None:
		params += [(u"--force", None, None)]
	if effective_snr_factor is not None:
		params += [(u"--effective-snr-factor", u"real_8", effective_snr_factor)]
	if vetoes_name is not None:
		params += [(u"--vetoes-name", u"lstring", vetoes_name)]
	if trigger_program is not None:
		params += [(u"--trigger-program", u"lstring", trigger_program)]
	if effective_snr is not None:
		params += [(u"--effective-snr", u"lstring", effective_snr)]
	if coinc_end_time_segment is not None:
		params += [(u"--coinc-end-time-segment", u"lstring", coinc_end_time_segment)]
	if exact_mass is not None:
		params += [(u"--exact-mass", None, None)]
	if verbose is not None:
		params += [(u"--verbose", None, None)]

	ligolw_process.append_process_params(xmldoc, process, params)

	return process


#
# =============================================================================
#
#                                     Main
#
# =============================================================================
#


#
# Command line
#


options, filenames = parse_command_line()


#
# Select effective SNR form
#


if options.effective_snr == "gstlal":
	ligolw_thinca.SnglInspiral.get_effective_snr = lambda self, fac: self.snr / math.sqrt(self.chisq)


#
# Select event_comparefunc form
#


if options.exact_mass:
	# Require exact mass coincidence
	event_comparefunc = ligolw_thinca.inspiral_coinc_compare_exact
else:
	event_comparefunc = ligolw_thinca.inspiral_coinc_compare


#
# Select ntuple_comparefunc form
#


if options.coinc_end_time_segs is not None:
	# Custom Ntuple Comparison Function
	def restricted_endtime_ntuple_comparefunc(events, offset_vector, seg = options.coinc_end_time_segs):
		"""
		Return False (ntuple should be retained) if the end time of the
		coinc is in the segmentlist segs.
		"""
		return ligolw_thinca.coinc_inspiral_end_time(events, offset_vector) not in seg

	ntuple_comparefunc = restricted_endtime_ntuple_comparefunc
else:
	ntuple_comparefunc = ligolw_thinca.default_ntuple_comparefunc


#
# Iterate over files.
#


for n, filename in enumerate(filenames):
	#
	# Load the file.
	#

	if options.verbose:
		print >>sys.stderr, "%d/%d:" % (n + 1, len(filenames)),
	xmldoc = ligolw_utils.load_filename(filename, verbose = options.verbose, contenthandler = ligolw.LIGOLWContentHandler)
	lsctables.table.InterningRowBuilder.strings.clear()

	#
	# Have we already processed it?
	#

	if ligolw_process.doc_includes_process(xmldoc, process_program_name):
		if options.verbose:
			print >>sys.stderr, "warning: %s already processed," % (filename or "stdin"),
		if not options.force:
			if options.verbose:
				print >>sys.stderr, "skipping"
			continue
		if options.verbose:
			print >>sys.stderr, "continuing by --force"

	#
	# Add an entry to the process table.
	#

	process = append_process(
		xmldoc,
		comment = options.comment,
		force = options.force,
		e_thinca_parameter = options.e_thinca_parameter,
		exact_mass = options.exact_mass,
		effective_snr_factor = options.effective_snr_factor,
		vetoes_name = options.vetoes_name,
		trigger_program = options.trigger_program,
		effective_snr = options.effective_snr,
		coinc_end_time_segment = options.coinc_end_time_segment,
		verbose = options.verbose
	)

	#
	# Hack the IDs.  LAL writes all triggers with event_id = 0.  This
	# value is used by other LAL programs to check if coincidence has
	# been performed yet so that behaviour cannot be changed but we
	# need the IDs to be unique before starting the coincidence engine.
	#
	# FIXME:  remove this when LAL writes trigger files with unique
	# IDs (unique within the output of lalapps_inspiral).
	#

	tbl = lsctables.SnglInspiralTable.get_table(xmldoc)
	tbl.set_next_id(lsctables.SnglInspiralID(0))
	for row in tbl:
		row.event_id = tbl.get_next_id()

	#
	# Extract veto segments if present.
	#
	# FIXME:  using the tools in the glue.ligolw.utils.segments module
	# it's not hard to modify the veto segments in the .xml to be just
	# those that intersect the search summary segments.  That way, if
	# multiple documents are inserted into the same database, or merged
	# with ligolw_add, the veto lists will not get duplicated.
	#

	if not ligolw_segments.has_segment_tables(xmldoc):
		if options.verbose:
			print >>sys.stderr, "warning: no segment definitions found, vetoes will not be applied"
		vetoes = None
	elif not ligolw_segments.has_segment_tables(xmldoc, name = options.vetoes_name):
		if options.verbose:
			print >>sys.stderr, "warning: document contains segment definitions but none named \"%s\", vetoes will not be applied" % options.vetoes_name
		vetoes = None
	else:
		vetoes = ligolw_segments.segmenttable_get_by_name(xmldoc, options.vetoes_name).coalesce()

	#
	# Run coincidence algorithm.
	#

	ligolw_thinca.ligolw_thinca(
		xmldoc,
		process_id = process.process_id,
		coinc_definer_row = ligolw_thinca.InspiralCoincDef,
		event_comparefunc = event_comparefunc,
		thresholds = options.e_thinca_parameter,
		ntuple_comparefunc = ntuple_comparefunc,
		effective_snr_factor = options.effective_snr_factor,
		veto_segments = vetoes,
		trigger_program = options.trigger_program,
		verbose = options.verbose
	)

	#
	# Close out the process table.
	#

	ligolw_process.set_process_end_time(process)

	#
	# Write back to disk, and clean up.
	#

	ligolw_utils.write_filename(xmldoc, filename, verbose = options.verbose, gz = (filename or "stdout").endswith(".gz"))
	xmldoc.unlink()
	lsctables.table.reset_next_ids(lsctables.TableByName.values())
