#!/usr/bin/python

# Copyright (C) 2015 Christopher M. Biwer
#
# 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.

import argparse
import h5py
import logging
import numpy
import pycbc.results
import sys
from glue import segments
from glue.ligolw import ligolw
from glue.ligolw import lsctables
from glue.ligolw import table
from glue.ligolw import utils
from pycbc.events.veto import get_segment_definer_comments
from pycbc.results import save_fig_with_metadata
from pycbc.workflow import SegFile
import pycbc.version

# parse command line
parser = argparse.ArgumentParser()
parser.add_argument("--version", action="version",
                        version=pycbc.version.git_verbose_msg)
parser.add_argument('--segment-files', type=str, nargs="+",
                        help='XML files with a segment definer table to read.')
parser.add_argument('--segment-names', type=str, nargs="+", required=False, default="",
                        help='Names of segments in the segment definer table.')
parser.add_argument('--description', type=str, required=False, default="",
                        help='Additional descriptive text for caption.')
parser.add_argument('--title-text', type=str, required=False,
                        help='Additional text to append to title.')
parser.add_argument('--output-file', type=str,
                        help='Path of the output HTML file.')
opts = parser.parse_args()

# setup log
logging.basicConfig(format='%(asctime)s:%(levelname)s : %(message)s',
                    level=logging.INFO,datefmt='%I:%M:%S')

# set column names
columns = (('Name', []),
           ('H1 Time (s)', []),
           ('L1 Time (s)', []),
           ('H1L1 Time (s)', []),
)
caption = "This table shows the cumulative amount of time for each segment. Shown are:"

# FIXME: set IFO list
ifos = ['H1', 'L1']

# loop over segment files from command line
seg_dict = {}
comment_dict = {}
for segment_file in opts.segment_files:

    # read segment definer table
    seg_dict.update(SegFile.from_segment_xml(segment_file).segment_dict)
    comment_dict.update(get_segment_definer_comments(open(segment_file, 'rb'),
                                                     include_version=False))

# loop over segment names
for segment_name in opts.segment_names:

    # allow user to find the and of segments if they use a "&" on the command line
    names = segment_name.split('&')

    # create a dict that will hold first segment_name in list for each IFO
    base_segs = segments.segmentlistdict({})

    # create a dict that will hold all other segment_name in list for each IFO
    comp_segs = segments.segmentlistdict({})

    # loop over IFOs
    for ifo in ifos:

        # make an empty list for each IFO
        base_segs[ifo] = segments.segmentlist([])
        comp_segs[ifo] = segments.segmentlist([])

        # loop over names from string split on "&"
        for name in names:

            # construct the first part of the segment flag
            # note that if you don't include a version on the command line then
            # it will select what ever version it finds first for that
            # segment flag
            segment_flag = ifo + ":" + name

            # loop over segments names from segment_definer table
            # from the segment XML files
            for key in seg_dict.keys():

                # if you find the segment flag for this IFO then add those
                # segments to the list and exit this for loop
                # this is a for loop because the version of the flag may not
                # be defined by the command line
                if key == segment_flag or key.startswith(segment_flag+':'):
                    if not len(base_segs[ifo]) and name == names[0]:
                        base_segs[ifo] = seg_dict[key].coalesce()
                    else:
                        comp_segs[ifo] = comp_segs[ifo].coalesce() + seg_dict[key].coalesce()
                        comp_segs[ifo].coalesce()
                    break

    # put comment in caption
    caption += " " + segment_name
    # FIXME: This only worked by fluke before. Not sure what it's supposed to
    #        do in the case where there are two segments anyway.
    #        Commenting for now
    #if comment_dict[key] != None:
    #    caption += " ("+comment_dict[key]+")"

    # get length of time of single-IFO segments in seconds
    if len(names) > 1:
        h1_len = abs( ( base_segs['H1'].coalesce() & comp_segs["H1"].coalesce() ).coalesce() )
        l1_len = abs( ( base_segs['L1'].coalesce() & comp_segs["L1"].coalesce() ).coalesce() )
    else:
        h1_len = abs( base_segs['H1'].coalesce() )
        l1_len = abs( base_segs['L1'].coalesce() )

    # find the AND of H1 and L1 time
    h1l1_base_segs = base_segs["H1"].coalesce() & base_segs["L1"].coalesce()
    h1l1_base_segs.coalesce()

    # find all the time that should overlap
    h1l1_comp_segs = comp_segs["H1"].coalesce() + comp_segs["L1"].coalesce()
    h1l1_comp_segs.coalesce()

    # get length of time in coincident H1L1 segments in seconds
    if len(names) > 1:
        h1l1_len = abs( (h1l1_base_segs & h1l1_comp_segs).coalesce() )
    else:
        h1l1_len = abs( h1l1_base_segs.coalesce() )

    # put values into columns
    columns[0][1].append(segment_name)
    columns[1][1].append(float(h1_len))
    columns[2][1].append(float(l1_len))
    columns[3][1].append(float(h1l1_len))

# cast columns into arrays
keys = [numpy.array(key, dtype=type(key[0])) for key,_ in columns]
vals = [numpy.array(val, dtype=type(val[0])) for _,val in columns]

# write HTML table
title = 'Segment Summary'
caption += '.'
if opts.title_text:
    title = title + ": " + opts.title_text
if opts.description:
   caption = caption + " " + opts.description

fig_kwds = {}
html_table = pycbc.results.table(vals, keys, page_size=10)
save_fig_with_metadata(str(html_table), opts.output_file,
                     fig_kwds=fig_kwds,
                     title=title,
                     cmd=' '.join(sys.argv),
                     caption=caption)
