#!/usr/bin/env python
#
# Copyright (C) 2010--2014  Kipp Cannon, Chad Hanna
#
# 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.
#
## @file gstlal_ll_inspiral_reset_likelihood
# A program to reset likelihood data after a burn in
#
# ### Command line interface
#
#		--verbose, action = "store_true", help = "Be verbose."
#		--likelihood-file", metavar = "filename", help = "Write likelihood files to disk and include the names in this cachefile file."
#		--zerolag-likelihood-file", metavar = "filename", help = "Write zerolag likelihood files to disk and include the names in this cachefile file."
#
#
# =============================================================================
#
#                                   Preamble
#
# =============================================================================
#


import sys
from gstlal import far
from glue.ligolw import utils as ligolw_utils
from glue.ligolw.utils import process as ligolw_process
from glue.ligolw import ligolw
from glue.lal import LIGOTimeGPS
from glue import segments
from optparse import OptionParser


def parse_command_line():
	parser = OptionParser(
		version = "Name: %%prog\n%s" % "" # FIXME
	)
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	parser.add_option("--background-ranking-file", metavar = "filename", help = "Write likelihood files to disk and include the names in this cachefile file.")
	parser.add_option("--zerolag-ranking-file", metavar = "filename", help = "Write zerolag likelihood files to disk and include the names in this cachefile file.")
	parser.add_option("--gps-start-time", metavar = "GPS", type = "int", help = "replace the live time segments with this GPS start time")
	parser.add_option("--gps-end-time", metavar = "GPS", type = "int", help = "replace the live time segments with this GPS end time")
	options, filenames = parser.parse_args()

	process_params = dict(options.__dict__)

	return options, process_params

options, process_params = parse_command_line()


like_cpd, like_rd, like_segs = far.parse_likelihood_control_doc(ligolw_utils.load_filename(options.background_ranking_file, verbose = options.verbose, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler))
_, zlike_rd, zlike_segs = far.parse_likelihood_control_doc(ligolw_utils.load_filename(options.zerolag_ranking_file, verbose = options.verbose, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler))

# replace thae live time with the user specified live time
new_like_segs = segments.segmentlistdict()
new_zlike_segs = segments.segmentlistdict()
for ifo in like_segs.keys():
	new_like_segs[ifo]  = segments.segmentlist([segments.segment([LIGOTimeGPS(float(options.gps_start_time)), LIGOTimeGPS(float(options.gps_end_time))])])
	new_zlike_segs[ifo] = segments.segmentlist([segments.segment([LIGOTimeGPS(float(options.gps_start_time)), LIGOTimeGPS(float(options.gps_end_time))])])

# Make the observed counts match the background sample but with the correct normalization
for instruments in zlike_rd.background_likelihood_rates:
	zlike_rd.zero_lag_likelihood_rates[instruments].array[:] = like_rd.background_likelihood_rates[instruments].array[:] / like_rd.background_likelihood_rates[instruments].array.sum() * zlike_rd.zero_lag_likelihood_rates[instruments].array.sum() / far.get_live_time(like_segs) * far.get_live_time(new_like_segs)
	# make the counts integers
	zlike_rd.zero_lag_likelihood_rates[instruments].array[zlike_rd.zero_lag_likelihood_rates[instruments].array < 1.] = 0.
	if options.verbose:
		print >> sys.stderr, "replacing %s with %d zero lag counts" % (str(instruments), zlike_rd.zero_lag_likelihood_rates[instruments].array.sum())
	if zlike_rd.zero_lag_likelihood_rates[instruments].array.sum() <= 100:
		raise ValueError("Must have greater than 100 counts in zero lag, have: %d" % zlike_rd.zero_lag_likelihood_rates[instruments].array.sum()) 


# write out the background file
xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())
process = ligolw_process.register_to_xmldoc(xmldoc, sys.argv[0], ifos = like_segs.keys(), paramdict = process_params)
far.gen_likelihood_control_doc(xmldoc, process, like_cpd, like_rd, new_like_segs, comment = u"reset background")
ligolw_utils.write_filename(xmldoc, options.background_ranking_file, gz = options.background_ranking_file.endswith(".gz"), verbose = options.verbose)

# write out the zerolag file
xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())
process = ligolw_process.register_to_xmldoc(xmldoc, sys.argv[0], ifos = zlike_segs.keys(), paramdict = process_params)
far.gen_likelihood_control_doc(xmldoc, process, None, zlike_rd, new_zlike_segs, comment = u"reset zerolag")
ligolw_utils.write_filename(xmldoc, options.zerolag_ranking_file, gz = options.zerolag_ranking_file.endswith(".gz"), verbose = options.verbose)
