#!/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
#
# =============================================================================
#

"""
Include the ranking data in the marginalized likelihood file
""" 

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

from glue import iterutils
from glue.text_progress_bar import ProgressBar
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 pylal import snglcoinc


#
# =============================================================================
#
#                                 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-file", metavar="FILENAME",
                        action="store", help = """The name of the likelihood
                        ratio data input file to use.""")
    parser.add_argument('--output-file', action="store", required=True,
                        metavar="OUTFILENAME", help="""The output file to write
                        the likelihood ratio file with ranking data 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 input document
#

in_xmldoc = ligolw_utils.load_url(args.likelihood_file, verbose = args.verbose, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler)

coincparamsdistributions, _, seglists = far.parse_likelihood_control_doc(in_xmldoc)
in_xmldoc.unlink()


coincparamsdistributions.horizon_distances = horizon_distances	# FIXME:  band-aid for coinc params func.  remove
likelihood_ratio_func = snglcoinc.LikelihoodRatio(coincparamsdistributions)


#
# dump combined distribution data file if requested.  we generate
# likelihood ratio histograms at this time, as well, even though we don't
# need them here because this helps parallelize that step in a full-scale
# analysis.  the program that reads these data files will need those
# histograms, and would take ages to build them itself, serially, as a
# single job.  We're about to spend a few hours processing database files
# and can generate those histograms in the background while we do
#


# take a moment to make sure we have SNR PDFs for all instrument combinations
for instruments in [instruments for n in range(2, len(seglists) + 1) for instruments in iterutils.choices(seglists.keys(), n)]:
    coincparamsdistributions.get_snr_joint_pdf(instruments, horizon_distances, progressbar = ProgressBar() if args.verbose else None)
# NOTE: This will use multiple cores (I think)!
ranking_data = far.RankingData(coincparamsdistributions, seglists.keys(), coincparamsdistributions.process_id, verbose=args.verbose)

#
# Collect the PDF writing thread if one was created and 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_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, ranking_data, 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)
