#!/home/travis/build/pycbc-build/environment/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 pycbc_glue import segments
from pycbc_glue.ligolw import ligolw
from pycbc_glue.ligolw import lsctables
from pycbc_glue.ligolw import table
from pycbc_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)
