#!/usr/bin/env python
#
# Copyright (C) 2010--2013  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.

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

"""
Combine gstlal likelihood files into one file and add in simulated injection
distributions.
"""

import pycbc
import pycbc.version
__author__  = "Kipp Cannon <kipp.cannon@ligo.org>, Ian Harry <ian.harry@ligo.org>"
__version__ = pycbc.version.git_verbose_msg
__date__    = pycbc.version.date
__program__ = "pycbc_combine_likelihood"

import sys
import pickle
import argparse
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

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


def parse_command_line():
    _desc = __doc__[1:]
    parser = argparse.ArgumentParser(description=_desc)
    parser.add_argument('--version', action='version', version=__version__)
    parser.add_argument('-v', '--verbose', action="store_true", default=False,
                        help="Be verbose.")
    parser.add_argument('--horizon-dist-file', action="store", required=True,
                        metavar="FILENAME", help="""The name of the pickle
                        file storing the horizon distances.""")
    parser.add_argument("-l", "--likelihood-urls", nargs='+', metavar="URL",
                        action="store", help = """The names of the likelihood
                        ratio data files to use. Filenames and URLs are
                        accepted.""")
    parser.add_argument("-s", "--synthesize-injections", metavar="N",
                        type=int, default=0, help="""Synthesize an injection
                        distribution with N injections, default 0""")
    parser.add_argument("-p", "--background-prior", metavar="N", default=0,
                        type=float, help="""Include an exponential background
                        prior with the maximum bin count = N, default 0, no
                        additional prior""")
    parser.add_argument('--output-file', action="store", required=True,
                        metavar="OUTFILENAME", help="""The output file to write
                        the combined likelihood ratio file to.""")
    args = parser.parse_args()
    return args

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


#
# command line
#


args = parse_command_line()

from gstlal import far

horizon_distances = pickle.load( open(args.horizon_dist_file, "rb"))


#
# load parameter distribution data
#


coincparamsdistributions = None
seglists = segments.segmentlistdict()
for n, likelihood_url in enumerate(args.likelihood_urls, start = 1):
    if args.verbose:
        print >>sys.stderr, "%d/%d:" % (n, len(args.likelihood_urls)),
    xmldoc = ligolw_utils.load_url(likelihood_url, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler, verbose = args.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 args.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())

# calculate injections before writing to disk
if args.synthesize_injections != 0:
    coincparamsdistributions.add_foreground_prior(n = args.synthesize_injections, segs = seglists, horizon_distances = horizon_distances, verbose = args.verbose)

# add a uniform prior to background, by default 0 is added so it has no effect
if args.background_prior != 0:
    coincparamsdistributions.add_background_prior(n = args.background_prior, instruments = seglists.keys(), verbose = args.verbose)

# compute the instrument combination counts
coincparamsdistributions.add_instrument_combination_counts(segs = seglists, verbose = args.verbose)

#
# rebuild event parameter PDFs (+= method has not constructed these
# correctly, and we might have added additional priors to the histograms),
# then initialize likeihood ratio evaluator
#

coincparamsdistributions.finish(verbose = args.verbose)

# Write output file
xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())
process = ligolw_process.register_to_xmldoc(xmldoc, u"gstlal_inspiral_calc_likelihood", 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, None, seglists)
ligolw_process.set_process_end_time(process)
ligolw_utils.write_filename(xmldoc, args.output_file, gz = (args.output_file or "stdout").endswith(".gz"), verbose = args.verbose)

