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

# parse command line
parser = argparse.ArgumentParser()
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(fromsegmentxml(open(segment_file, 'rb'), return_dict=True))
    comment_dict.update(get_segment_definer_comments(open(segment_file, 'rb')))

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

    # allow user to find the and of two segments if they use a '&' on the command line
    names = segment_name.split('&')
    if len(names) > 2:
        logging.info('You can only compare two segments but was provided %s', segment_name)
        continue

    # get segments for each IFO
    segs = segments.segmentlistdict()
    for ifo in ifos:
        for key in seg_dict.keys():
            if key.startswith(ifo+':'+names[0]):
                segs[ifo] = seg_dict[key]
    if not len(segs.keys()):
        logging.info('Did not find a segment definition for %s', segment_name)
        continue

    # get second segments for each IFO if comparing two segments
    if len(names) == 2:
        segs_tmp = segments.segmentlistdict()
        for ifo in ifos:
            for key in seg_dict.keys():
                if key.startswith(ifo+':'+names[1]):
                    segs_tmp[ifo] = seg_dict[key]
        if not len(segs.keys()):
            logging.info('Did not find a segment definition for %s', segment_name)
            continue

        # find overlap
        segs = segs.coalesce() & segs_tmp.coalesce()

    # put comment in caption
    caption += " " + segment_name
    if comment_dict[key] != None:
        caption += " ("+comment_dict[key]+")"

    # get length of time of segments in seconds
    h1_len = abs(segs['H1'])
    l1_len = abs(segs['L1'])
    h1l1_len = abs( segs['H1'] & segs['L1'] )

    # put values into columns
    columns[0][1].append(segment_name)
    columns[1][1].append(h1_len)
    columns[2][1].append(l1_len)
    columns[3][1].append(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)
