#!/usr/bin/env python

# Copyright (C) 2015 Ian Harry
#
# 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 3 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.

"""
Take a coinc xml file containing multiple events and upload to gracedb.
"""

import os
import argparse
import logging
from ligo.gracedb.rest import GraceDb
import numpy as np
import h5py

import pycbc
import lal
import lal.series
from glue.ligolw import ligolw
from glue.ligolw import table
from glue.ligolw import lsctables
from glue.ligolw import utils as ligolw_utils
from ligo.segments import segment, segmentlist
from pycbc.io.live import gracedb_tag_with_version

def check_gracedb_for_event(gdb_handle, query, far):
    """
    Check if there is an event in gracedb with the queried string
    which matches the FAR given
    """
    gdb_events_match_query = list(gdb_handle.events(query=query))
    ifar = 1. / lal.YRJUL_SI / far
    for gdb_event in gdb_events_match_query:
        # Test each gracedb event to see if the FAR matches this event
        if np.abs(gdb_event['far'] - far) < 1e-16:
            # If an event has been found which matches the FAR
            logging.info('Event already exists in GraceDb server with '
                         'time %.3f and IFAR %.3e: %s',
                         gdb_event['gpstime'], ifar,
                         gdb_event['graceid'])
            return True

    # If no event has been found, log this, and return False
    logging.info('No event found in GraceDb with IFAR %.3e when using '
                 'query: "%s"', ifar, query)
    return False

class LIGOLWContentHandler(ligolw.LIGOLWContentHandler):
    pass

logging.basicConfig(format='%(asctime)s:%(levelname)s : %(message)s',
                    level=logging.INFO)

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--psd-files", nargs='+', required=True,
                    help='HDF file(s) containing the PSDs to upload')
parser.add_argument('--input-file', required=True, type=str,
                    help='Input LIGOLW XML file of coincidences.')
parser.add_argument('--log-message', type=str, metavar='MESSAGE',
                    help='Add a log entry to each upload with the given message')
parser.add_argument('--testing', action="store_true", default=False,
                    help="Upload event to the TEST group of gracedb.")
parser.add_argument('--min-ifar', type=float, metavar='YEARS',
                    help='Only upload events more significant than given IFAR')
parser.add_argument('--production-server', action="store_true", default=False,
                    help="Upload event to production graceDB. If not given "
                         "events will be uploaded to playground server.")
parser.add_argument('--force-overwrite', action='store_true', default=False,
                    help="GraceDb instance will be checked for if an event "
                         "with the same event time and FAR already exist. "
                         "If so, event will not be uploaded")
parser.add_argument('--query-string', default='pycbc',
                    help="If not using --force-overwrite, add a string to "
                         "gracedb query to further filter events. "
                         "Default='pycbc'")
parser.add_argument('--search-id-string', default='AllSky',
                    help="Using T050017-v1 naming convention for "
                         "output XML filename. This is a string for "
                         "search identifier in filename, e.g. 'AllSky' "
                         "would give H1L1V1-PYCBC_AllSky-1234567890-1.xml. "
                         "See https://dcc.ligo.org/LIGO-T050017/public. "
                         "Default: 'AllSky'")
parser.add_argument('--output-directory', default=os.getcwd(),
                    help="Output directory for locally stored XML and PSD "
                         "files. Default: current directory")
parser.add_argument('--no-upload', action='store_true',
                    help="Flag used to indicate that we are not uploading to "
                         "GraceDb.")

args = parser.parse_args()

if args.production_server:
    gracedb = GraceDb()
else:
    gracedb = GraceDb(service_url='https://gracedb-playground.ligo.org/api/')

lsctables.use_in(LIGOLWContentHandler)

xmldoc = ligolw_utils.load_filename(args.input_file,
                             contenthandler=LIGOLWContentHandler)

class psd_segment(segment):
    def __new__(cls, psd, *args):
        return segment.__new__(cls, *args)
    def __init__(self, psd, *args):
        self.psd = psd

psds = {}
for psd_file in args.psd_files:
    (ifo, group), = h5py.File(psd_file, "r").items()
    psd = [group["psds"][str(i)] for i in range(len(group["psds"].keys()))]
    psds[ifo] = segmentlist(psd_segment(*segargs) for segargs in zip(
        psd, group["start_time"], group["end_time"]))

coinc_table = table.get_table(xmldoc, lsctables.CoincTable.tableName)
coinc_inspiral_table = table.get_table(xmldoc,
                                       lsctables.CoincInspiralTable.tableName)
coinc_event_map_table = table.get_table(xmldoc,
                                        lsctables.CoincMapTable.tableName)

sngl_inspiral_table = table.get_table(xmldoc,
                                      lsctables.SnglInspiralTable.tableName)

xmldoc.childNodes[-1].removeChild(sngl_inspiral_table)
xmldoc.childNodes[-1].removeChild(coinc_event_map_table)
xmldoc.childNodes[-1].removeChild(coinc_inspiral_table)
xmldoc.childNodes[-1].removeChild(coinc_table)

for event in coinc_table:
    coinc_event_table_curr = lsctables.New(lsctables.CoincTable)
    coinc_event_table_curr.append(event)
    coinc_inspiral_table_curr = lsctables.New(lsctables.CoincInspiralTable)
    coinc_event_map_table_curr = lsctables.New(lsctables.CoincMapTable)
    sngl_inspiral_table_curr = lsctables.New(lsctables.SnglInspiralTable)

    coinc_event_id = event.coinc_event_id
    for coinc_insp in coinc_inspiral_table:
        if coinc_insp.coinc_event_id == event.coinc_event_id:
            coinc_inspiral_table_curr.append(coinc_insp)

    if args.min_ifar is not None and \
            coinc_inspiral_table_curr[0].combined_far > 1./args.min_ifar/lal.YRJUL_SI:
        continue

    time = coinc_inspiral_table_curr[0].end_time
    if not args.force_overwrite:
        far = coinc_inspiral_table_curr[0].combined_far
        query = args.query_string + ' %.3f .. %.3f' % (time - 1, time + 1)
        if check_gracedb_for_event(gracedb, query, far):
            continue

    sngl_ids = []
    for coinc_map in coinc_event_map_table:
        if coinc_map.coinc_event_id == event.coinc_event_id:
            coinc_event_map_table_curr.append(coinc_map)
            sngl_ids.append(coinc_map.event_id)

    psddict = {}
    for sngl in sngl_inspiral_table:
        if sngl.event_id in sngl_ids:
            sngl_inspiral_table_curr.append(sngl)

            try:
                psd = psds[sngl.ifo]
            except KeyError:
                parser.error(
                    "--psd-files {0}: no PSDs found for detector {1}".format(
                    " ".join(args.psd_files), sngl.ifo))

            try:
                psd = psd[psd.find(sngl.get_end())].psd
            except ValueError:
                parser.error(
                    "--psd-files {0}: no PSD found for detector {1} "
                    "at GPS time {2}".format(
                    " ".join(args.psd_files), sngl.ifo, sngl.get_end()))

            flow = psd.file.attrs['low_frequency_cutoff']
            df = psd.attrs["delta_f"]
            kmin = int(flow / df)

            fseries = lal.CreateREAL8FrequencySeries(
                "psd", psd.attrs["epoch"], kmin * df, df,
                lal.StrainUnit**2 / lal.HertzUnit, len(psd) - kmin)
            fseries.data.data = psd[kmin:] / np.square(pycbc.DYN_RANGE_FAC)
            psddict[sngl.ifo] = fseries

    xmldoc.childNodes[-1].appendChild(coinc_event_table_curr)
    xmldoc.childNodes[-1].appendChild(coinc_inspiral_table_curr)
    xmldoc.childNodes[-1].appendChild(coinc_event_map_table_curr)
    xmldoc.childNodes[-1].appendChild(sngl_inspiral_table_curr)

    ifos = sorted([sngl.ifo for sngl in sngl_inspiral_table_curr])
    ifos_str = ''.join(ifos)
    id_str = args.search_id_string
    filename_xml = "{}-PYCBC_{}-{:d}-1.xml".format(ifos_str, id_str, time)
    filename_psd = "{}-PYCBC_{}-{:d}-1_psd.xml.gz".format(ifos_str, id_str, time)
    fullpath_xml = os.path.join(args.output_directory, filename_xml)
    fullpath_psd = os.path.join(args.output_directory, filename_psd)

    ligolw_utils.write_filename(xmldoc, fullpath_xml)
    psd_xmldoc = lal.series.make_psd_xmldoc(psddict)
    ligolw_utils.write_filename(psd_xmldoc, fullpath_psd, gz=True)

    if args.no_upload:
        # We have saved the file, and aren't uploading, so reset the tables
        # and move to the next event
        logging.info("Not uploading event")
        xmldoc.childNodes[-1].removeChild(coinc_event_table_curr)
        xmldoc.childNodes[-1].removeChild(coinc_inspiral_table_curr)
        xmldoc.childNodes[-1].removeChild(coinc_event_map_table_curr)
        xmldoc.childNodes[-1].removeChild(sngl_inspiral_table_curr)
        continue

    group_tag = 'Test' if args.testing else 'CBC'
    r = gracedb.createEvent(group_tag, 'pycbc', filename_xml,
                            filecontents=open(fullpath_xml, "rb").read(),
                            search=id_str, offline=True).json()

    logging.info("Uploaded event %s.", r["graceid"])

    # upload PSD
    gracedb.writeLog(
        r["graceid"], "PyCBC PSD estimate from the time of event",
        "psd.xml.gz", open(fullpath_psd, "rb").read(), "psd").json()
    logging.info("Uploaded file %s to event %s.", filename_psd, r["graceid"])

    # add info for tracking code version
    gracedb_tag_with_version(gracedb, r['graceid'])

    # document the absolute path to the input file
    input_file_str = 'Candidate uploaded from ' \
        + os.path.abspath(args.input_file)
    gracedb.writeLog(r['graceid'], input_file_str)

    # add the custom log message, if provided
    if args.log_message is not None:
        gracedb.writeLog(r['graceid'], args.log_message,
                         tag_name=['analyst_comments'])

    xmldoc.childNodes[-1].removeChild(coinc_event_table_curr)
    xmldoc.childNodes[-1].removeChild(coinc_inspiral_table_curr)
    xmldoc.childNodes[-1].removeChild(coinc_event_map_table_curr)
    xmldoc.childNodes[-1].removeChild(sngl_inspiral_table_curr)