#!/usr/bin/python
#
# Copyright (C) 2006--2009,2013,2014  Kipp Cannon
#
# 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
#
# =============================================================================
#


"""
Command-line interface to CBC injection identification code. Accepts as input
one LIGO Light Weight XML file, containing CBC candidates and a
list of injections, and adds entries to the coincidence tables indicating which
events match which injections
"""


import argparse
import sys


from glue.ligolw import lsctables
from glue.ligolw import utils as ligolw_utils
from glue.ligolw.utils import process as ligolw_process
from pylal import ligolw_inspinjfind
from pylal import SnglInspiralUtils
import pycbc.version

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


#
# =============================================================================
#
#                                 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("-i", "--input-file", action="store", type=str,
                        default=None, required=True,
                        help="""Input XML file to read. Can only
                                input one at a time.""")
    parser.add_argument("-o", "--output-file", action="store", type=str,
                        default=None, required=True,
                        help="Name of output XMl file to save to.")

    parser.add_argument("-f", "--force", action="store_true",
                        help="Process even if file has already been processed.")
    parser.add_argument("--comment", metavar="text", 
                        help="""Set the comment string to be written to the
                                process table (default = None).""")
    parser.add_argument("-c", "--match-algorithm", metavar="[inspiral]",
                        default="inspiral",
                        help="""Set the algorithm used to match event candidates
                                with injections (default = \"inspiral\").""")
    parser.add_argument("-w", "--time-window", metavar="seconds", type=float,
                        default=9.,
                        help="""Set the time window used by the \"inspiral\"
                                match algorithm (default = 9.0).""")
    parser.add_argument("-v", "--verbose", action="store_true",
                        help="Be verbose.")
    args = parser.parse_args()
    if args.match_algorithm not in ("inspiral",):
        raise ValueError("unrecognized --match-algorithm \"%s\"" \
                         % args.match_algorithm)
    if args.time_window < 0.:
        raise ValueError("--time-window must be non-negative")

    return args


#
# =============================================================================
#
#                                    Input
#
# =============================================================================
#


#
# Use interning row builder to save memory.
#


lsctables.table.RowBuilder = lsctables.table.InterningRowBuilder


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


#
# command line
#


args = parse_command_line()

# must match columns in sngl_inspiral table
search = {"inspiral": "inspiral"}[args.match_algorithm]

#
# set compare functions
#

# FIXME: Ian is confused? Aren't these the same??
def InspiralSnglCompare(sim, inspiral,
                  twindow=ligolw_inspinjfind.LIGOTimeGPS(args.time_window)):
    return SnglInspiralUtils.CompareSnglInspiral(sim, inspiral,twindow=twindow)

def NearCoincCompare(sim, inspiral,
                  twindow=ligolw_inspinjfind.LIGOTimeGPS(args.time_window)):
    return SnglInspiralUtils.CompareSnglInspiral(sim, inspiral, twindow=twindow)


snglcomparefunc = {"inspiral": InspiralSnglCompare}[args.match_algorithm]
nearcoinccomparefunc = {"inspiral": NearCoincCompare}[args.match_algorithm]

#
# load the document
#

xmldoc = ligolw_utils.load_filename(args.input_file,
                        contenthandler=ligolw_inspinjfind.LIGOLWContentHandler,
                        verbose = args.verbose)
lsctables.table.InterningRowBuilder.strings.clear()

#
# add process metadata to document
#

process = ligolw_inspinjfind.append_process(xmldoc, 
                                          match_algorithm=args.match_algorithm,
                                          comment=args.comment)

#
# run inspinjfind algorithm
#

ligolw_inspinjfind.ligolw_inspinjfind(xmldoc, process, search, snglcomparefunc,
                                nearcoinccomparefunc,
                                end_time_bisect_window=1.5*args.time_window,
                                verbose = args.verbose)

#
# close out the process metadata
#

ligolw_process.set_process_end_time(process)

#
# done
#

ligolw_utils.write_filename(xmldoc, args.output_file, verbose=args.verbose,
                                         gz = args.output_file.endswith(".gz"))