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

"""
Calculate likelihood value for coincident triggers in input 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_calculate_likelihood"


import argparse
import sys
import pickle
import sqlite3

from glue.ligolw import ligolw
from glue.ligolw import dbtables
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 search_summary as ligolw_search_summary
from glue.ligolw.utils import segments as ligolw_segments
from glue import segments
from pylal import ligolw_burca2
from pylal import ligolw_thinca
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("-l", "--likelihood-file", metavar="FILENAME",
                        action="store", help = """The name of the likelihood
                        ratio data input file to use.""")
    parser.add_argument("-T", "--trigger-file", metavar="FILENAME",
                        action="store", help = """The name of the file
                        containing the coincident triggers to calculate the
                        likelihood for.""")
    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('--output-file', action="store", required=True,
                        metavar="OUTFILENAME", help="""The output file to write
                        the likelihood ratio file with ranking data to.""")
    parser.add_argument('-t', '--tmp-space', action="store", default=None,
                    metavar='PATH', help="""Path to a directory suitable for use
                    as a work area while manipulating the database file.  The
                    database file will be worked on in this directory, and then
                    moved to the final location when complete.  This option is
                    intended to improve performance when running in a networked
                    environment, where there might be a local disk with higher
                    bandwidth than is available to the filesystem on which the
                    final output will reside.""")
    parser.add_argument("--vetoes-name", metavar="name",
                        help="""Set the name of the segment lists to use as
                                vetoes (default = do not apply vetoes).""")
    parser.add_argument("--trim-database", action="store_true", 
                        help="""Delete events that are found to be below the pipeline's likelihood ratio threshold (default = do not delete them).  Deleting the events saves a significant amount of disk space but is inconvenient during pipeline development and tuning as it makes it impossible to rerank the events later with a different ranking statistic.""")
    parser.add_argument("--drop-veto-info", action = "store_true", help = "Remove the segment, segment_definer, veto_definer and segment_summary tables from the database as well as the associated rows in the process and process_params tables.")


    args = parser.parse_args()
    return args


#
# =============================================================================
#
#                   Support Funcs for Likelihood Ratio Code
#
# =============================================================================
#


def sngl_inspiral_events_func(cursor, coinc_event_id, row_from_cols):
	return map(row_from_cols, cursor.execute("""
SELECT
	sngl_inspiral.*
FROM
	sngl_inspiral
	JOIN coinc_event_map ON (
		coinc_event_map.table_name == 'sngl_inspiral'
		AND coinc_event_map.event_id == sngl_inspiral.event_id
	)
WHERE
	coinc_event_map.coinc_event_id == ?
	""", (coinc_event_id,)))


def sngl_inspiral_veto_func(event, vetoseglists):
	# return True if event should be *retained*
	return event.ifo not in vetoseglists or event.get_end() not in vetoseglists[event.ifo]


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


#
# command line
#

args = parse_command_line()

from gstlal import far

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

#
# load input likelihood 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)

#
# Open the database file.
#

working_filename = dbtables.get_connection_filename(args.trigger_file,
                                 tmp_path=args.tmp_space, verbose=args.verbose)
connection = sqlite3.connect(working_filename)
if args.tmp_space is not None:
    dbtables.set_temp_store_directory(connection, args.tmp_space,
                                                          verbose=args.verbose)
xmldoc = dbtables.get_xml(connection)

#
# Summarize the database, and record our passage.
#

try:
    coinc_def_id = lsctables.CoincDefTable.get_table(xmldoc).get_coinc_def_id(ligolw_thinca.InspiralCoincDef.search, ligolw_thinca.InspiralCoincDef.search_coinc_type, create_new = False)
except KeyError:
    if args.verbose:
        print >>sys.stderr, "document does not contain inspiral coincidences."
    coinc_def_id = None

process = ligolw_process.register_to_xmldoc(xmldoc, u"gstlal_inspiral_calc_likelihood", {})
search_summary = ligolw_search_summary.append_search_summary(xmldoc, process, ifos = seglists.keys(), inseg = seglists.extent_all(), outseg = seglists.extent_all())

sngl_inspiral_table = lsctables.SnglInspiralTable.get_table(xmldoc)

offset_vectors = lsctables.TimeSlideTable.get_table(xmldoc).as_dict()

if args.vetoes_name is not None:
    vetoseglists = ligolw_segments.segmenttable_get_by_name(xmldoc, args.vetoes_name).coalesce()
else:
    vetoseglists = segments.segmentlistdict()

#
# Run likelihood ratio calculation.
#

if coinc_def_id is not None:
    ligolw_burca2.assign_likelihood_ratios(
            connection = connection,
            coinc_def_id = coinc_def_id,
            offset_vectors = offset_vectors,
            vetoseglists = vetoseglists,
            events_func = lambda cursor, coinc_event_id: sngl_inspiral_events_func(cursor, coinc_event_id, sngl_inspiral_table.row_from_cols),
            veto_func = sngl_inspiral_veto_func,
            likelihood_ratio_func = likelihood_ratio_func,
            likelihood_params_func = coincparamsdistributions.coinc_params,
            verbose = args.verbose
)

#
# Delete low significance events to reduce database size
#

if args.trim_database and (coinc_def_id is not None):
    cursor = connection.cursor()
    cursor.execute("""DELETE FROM coinc_event WHERE coinc_def_id == ? AND likelihood < ?;""", (coinc_def_id, far.RankingData.likelihood_ratio_threshold))
    cursor.execute("""DELETE FROM coinc_inspiral WHERE coinc_event_id NOT IN (SELECT coinc_event_id FROM coinc_event);""")
    cursor.execute("""DELETE FROM coinc_event_map WHERE coinc_event_id NOT IN (SELECT coinc_event_id FROM coinc_event);""")
   # FIXME:  don't hard-code parameter name
    cursor.execute("""DELETE FROM sngl_inspiral WHERE snr < (SELECT value FROM process_params WHERE process_params.process_id == sngl_inspiral.process_id AND param == "--singles-threshold") AND event_id NOT IN (SELECT event_id FROM coinc_event_map WHERE table_name == "sngl_inspiral");""")
    cursor.execute("""VACUUM;""")
    cursor.close()

if args.drop_veto_info:
    cursor = connection.cursor()

    query_text = """
    DELETE FROM process
        WHERE process_id IN (
            SELECT DISTINCT process_id
            FROM segment_definer);"""
    cursor.execute(query_text)
    query_text = """
    DELETE FROM process_params
        WHERE process_id IN (
            SELECT DISTINCT process_id
            FROM segment_definer);"""
    cursor.execute(query_text)
    cursor.execute("DROP TABLE segment;")
    cursor.execute("DROP TABLE segment_summary;")
    cursor.execute("DROP TABLE segment_definer;")
    if args.verbose:
        print >> sys.stderr, "segment, segment-definer & segment-summary tables dropped "
    query_text = """
    DELETE FROM process
        WHERE process_id IN (
            SELECT DISTINCT process_id
            FROM veto_definer);"""
    cursor.execute(query_text)
    query_text = """
    DELETE FROM process_params
        WHERE process_id IN (
            SELECT DISTINCT process_id
            FROM veto_definer);"""
    cursor.execute(query_text)

    # remove veto_definer table from xmldoc
    query_text = "DROP TABLE veto_definer;" 
    cursor.execute(query_text)
    if args.verbose:
        print >> sys.stderr, "Veto-Definer table dropped"

#
# Close out process metadata
#

ligolw_process.set_process_end_time(process)
connection.cursor().execute("UPDATE process SET end_time = ? WHERE process_id == ?", (process.end_time, process.process_id))

#
# Clean up.
#

xmldoc.unlink()
connection.commit()
connection.close()
dbtables.put_connection_filename(args.output_file, working_filename, verbose=args.verbose)

