#!/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
# A program to compute the noise probability distributions of likehood ratios for inspiral triggers
#
# ### Command line interface
#
#	+ `--likelihood-cache` [filename]: Also load the likelihood ratio data files listsed in this LAL cache.  See lalapps_path2cache for information on how to produce a LAL cache file.
#	+ `--verbose`: Be verbose.
#	+ `--ranking-stat-samples` [N] (int): Construct ranking statistic histograms by drawing this many samples from the ranking statistic generator (default = 1000000).
#	+ `--output` [filename]: Write merged raw likelihood data and likelihood ratio histograms to this LIGO Light-Weight XML file.

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


from optparse import OptionParser
import sys


from glue import iterutils
from glue.lal import CacheEntry
from glue.ligolw import ligolw
from glue.ligolw import utils as ligolw_utils
from glue.ligolw.utils import process as ligolw_process
from glue.ligolw.utils import search_summary as ligolw_search_summary
from glue import segments
from gstlal import far


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


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


def parse_command_line():
	parser = OptionParser(
		version = "Name: %%prog\n%s" % "" # FIXME
	)
	parser.add_option("--likelihood-cache", metavar = "filename", help = "Also load the likelihood ratio data files listsed in this LAL cache.  See lalapps_path2cache for information on how to produce a LAL cache file.")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	parser.add_option("--ranking-stat-samples", metavar = "N", default = 10000000, type = "int", help = "Construct ranking statistic histograms by drawing this many samples from the ranking statistic generator (default = 10000000).")
	parser.add_option("--samples-file", metavar = "filename", action = "append", help = "Load the results of previous sampler runs from this file, and add the samples we generate to it.  Can be given multiple times.")
	parser.add_option("--samples-cache", metavar = "filename", help = "Load the results of previous sampler runs from the files listed in this LAL cache.  See lalapps_path2cache for information on how to produce a LAL cache file.")
	parser.add_option("--output", metavar = "filename", help = "Write merged raw likelihood data and likelihood ratio histograms to this LIGO Light-Weight XML file.")
	options, urls = parser.parse_args()

	paramdict = options.__dict__.copy()

	if options.likelihood_cache is not None:
		urls += [CacheEntry(line).url for line in open(options.likelihood_cache)]
	if not urls:
		raise ValueError("must provide some likelihood files")

	if options.samples_file is None:
		options.samples_file = []
	if options.samples_cache is not None:
		options.samples_file += [CacheEntry(line).url for line in open(options.samples_cache)]

	if options.output is None:
		raise ValueError("must set --output")

	return options, urls, paramdict


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


#
# command line
#


options, urls, paramdict = parse_command_line()


#
# load parameter distribution data and, optionally, samples from previous
# runs
#


old_ranking_data = None
for n, filename in enumerate(options.samples_file, start = 1):
	if options.verbose:
		print >>sys.stderr, "%d/%d:" % (n, len(options.samples_file)),
	xmldoc = ligolw_utils.load_url(filename, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler, verbose = options.verbose)
	ignored, this_ranking_data, ignored = far.parse_likelihood_control_doc(xmldoc)
	xmldoc.unlink()
	if this_ranking_data is None:
		raise ValueError("%s does not contain ranking statistic distribution data" % filename)
	if old_ranking_data is None:
		old_ranking_data = this_ranking_data
	else:
		old_ranking_data += this_ranking_data

coincparamsdistributions = None
seglists = segments.segmentlistdict()
for n, likelihood_url in enumerate(urls, start = 1):
	if options.verbose:
		print >>sys.stderr, "%d/%d:" % (n, len(urls)),
	xmldoc = ligolw_utils.load_url(likelihood_url, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler, verbose = options.verbose)
	this_coincparamsdistributions, ignored, this_seglists = far.parse_likelihood_control_doc(xmldoc)
	xmldoc.unlink()
	if this_coincparamsdistributions is None:
		raise ValueError("%s does not contain parameter distribution data" % likelihood_url)
	if coincparamsdistributions is None:
		coincparamsdistributions = this_coincparamsdistributions
	else:
		coincparamsdistributions += this_coincparamsdistributions
	seglists |= this_seglists

if options.verbose:
	print >>sys.stderr, "total livetime:\n\t%s" % ",\n\t".join("%s = %s s" % (instrument, str(abs(segs))) for instrument, segs in seglists.items())

# Compute the probability of instruments given signal
coincparamsdistributions.populate_prob_of_instruments_given_signal(segs = seglists, n = 1.0, verbose = options.verbose)

# Compute the intrument combination counts
coincparamsdistributions.add_instrument_combination_counts(segs = seglists, verbose = options.verbose)


#
# rebuild event parameter PDFs (+= method has not constructed these
# correctly)
#


coincparamsdistributions.finish(verbose = options.verbose)
if old_ranking_data is not None:
	old_ranking_data.finish(verbose = options.verbose)


#
# take a moment to make sure we have SNR PDFs for all instrument
# combinations so that we don't end up doing this inside the sampling
# threads
#


all_instrument_combos = [instruments for n in range(2, len(seglists) + 1) for instruments in iterutils.choices(seglists.keys(), n)]
for horizon_distances in coincparamsdistributions.horizon_history.all():
	for instruments in all_instrument_combos:
		coincparamsdistributions.get_snr_joint_pdf(instruments, horizon_distances, verbose = options.verbose)


#
# generate likelihood ratio histograms
#


ranking_data = far.RankingData(coincparamsdistributions, instruments = seglists.keys(), nsamples = options.ranking_stat_samples, verbose = options.verbose)
if old_ranking_data is not None:
	ranking_data += old_ranking_data


#
# Write the parameter and ranking statistic distribution data to a file
#


xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())
process = ligolw_process.register_to_xmldoc(xmldoc, u"gstlal_inspiral_calc_rank_pdfs", paramdict = paramdict)
search_summary = ligolw_search_summary.append_search_summary(xmldoc, process, ifos = seglists.keys(), inseg = seglists.extent_all(), outseg = seglists.extent_all())
far.gen_likelihood_control_doc(xmldoc, process, coincparamsdistributions, ranking_data, seglists)
ligolw_process.set_process_end_time(process)
ligolw_utils.write_filename(xmldoc, options.output, gz = (options.output or "stdout").endswith(".gz"), verbose = options.verbose)
