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

"""
Get gstlal horizon distance information from a set of input files
"""

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


import argparse
import sqlite3
import sys
import pickle

from glue.ligolw import dbtables
from glue.ligolw import lsctables

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


def parse_command_line():
    """
    Parse the command line, return options and check for consistency among the
    options.
    """
    _desc = __doc__[1:]
    parser = argparse.ArgumentParser(description=_desc)
    parser.add_argument('--version', action='version', version=__version__)
    parser.add_argument('input_files', nargs='+', action="store",
                        metavar="INPUT_FILE", help="""List of files to get the
                        horizon distances from.""")
    parser.add_argument('--output-file', action="store", required=True,
                        metavar="OUTFILENAME", help="""The name of the pickle
                        file to dump this information to.""")
    parser.add_argument('-v', '--verbose', action="store_true", default=False,
                        help="Be verbose.")

    args = parser.parse_args()
    return args

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


#
# command line
#


args = parse_command_line()

from gstlal import inspiral

#
# retrieve horizon distance data
#


if args.verbose:
	print >>sys.stderr, "retrieving horizon distances ..."

horizon_distances = {}
for n, filename in enumerate(args.input_files, 1):
    #
    # open the database file, do the query in-place
    #

    if args.verbose:
        print >>sys.stderr, "%d/%d: %s" % (n, len(args.input_files), filename)

    working_filename = dbtables.get_connection_filename(filename, tmp_path = None, verbose = args.verbose)
    connection = sqlite3.connect(working_filename)

    #
    # retrieve horizon distances
    #

    for instruments, start_time, start_time_ns, name, dist in connection.cursor().execute("""
SELECT
	summ_value.ifo,
	summ_value.start_time,
	summ_value.start_time_ns,
	summ_value.name,
	summ_value.value
FROM
	summ_value
WHERE
	program == "gstlal_inspiral"
	AND name LIKE "Dh:%";
	"""):
        # FIXME:  check that these are what we expect?
        m1, m2, snr_threshold = inspiral.CoincsDocument.summ_value_name_decode(name)
        for instrument in lsctables.instrument_set_from_ifos(instruments):
            horizon_distances.setdefault(instrument, []).append((lsctables.LIGOTimeGPS(start_time, start_time_ns), dist))

     # Also for pycbc

    for instruments, start_time, start_time_ns, name, dist, comment in connection.cursor().execute("""
SELECT
        summ_value.ifo,
        summ_value.start_time,
        summ_value.start_time_ns,
        summ_value.name,
        summ_value.value,
        summ_value.comment
FROM
        summ_value
WHERE
        program == "PyCBC-INSPIRAL"
        AND name == "inspiral_effective_distance";
        """):
        # FIXME:  check that these are what we expect?
        m1, m2, snr_threshold = map(float, comment.split("_"))
        for instrument in lsctables.instrument_set_from_ifos(instruments):
            horizon_distances.setdefault(instrument, []).append((lsctables.LIGOTimeGPS(start_time, start_time_ns), dist))


    connection.close()
    dbtables.discard_connection_filename(filename, working_filename, verbose = args.verbose)

#
# put horizon distances for each instrument in time order
#

for dists in horizon_distances.values():
    dists.sort()

# FIXME:  for now we just take the average of the distances recorded for
# each instrument.  this needs to be worked out properly so that events are
# ranked with the appropriate distance
horizon_distances = dict((instrument, sum(dist for t, dist in dists) / len(dists)) for instrument, dists in horizon_distances.items())

pickle.dump( horizon_distances, open( args.output_file, "wb" ) )
