#!/bin/env python
# Copyright (C) 2015 Alexander Harvey Nitz
#
# 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.
""" Make tables describing a coincident foreground event"""
import h5py, argparse, logging, sys
import matplotlib; matplotlib.use('Agg')
import numpy, lal, datetime
import pycbc.version, pycbc.events, pycbc.results, pycbc.pnutils

parser = argparse.ArgumentParser()
parser.add_argument('--version', action='version',
    version=pycbc.version.git_verbose_msg)
parser.add_argument('--verbose', action='store_true')
parser.add_argument('--single-trigger-files', nargs='+', 
    help="HDF format single detector trigger files for the full data run")
parser.add_argument('--bank-file',
    help="HDF format template bank file")
parser.add_argument('--output-file')
parser.add_argument('--statmap-file', required=True,
    help="HDF format clustered coincident statmap file containing the result "
         "triggers. Required")
parser.add_argument('--statmap-file-subspace-name', default='foreground',
    help="If given look in this 'sub-directory' of the HDF file for triggers, "
         "takes a default value of 'foreground'.")
trig_input = parser.add_mutually_exclusive_group(required=True)
trig_input.add_argument('--n-loudest', type=int,
    help="Examine the n'th loudest trigger, use with statmap file")
trig_input.add_argument('--trigger-id', type=int,
    help="Examine the trigger with specified ID, use with statmap file. An "
         "alternative to --n-loudest. Cannot be used together")

args = parser.parse_args()
pycbc.init_logging(args.verbose)

# Get the nth loudest trigger from the output of pycbc_coinc_statmap
f = h5py.File(args.statmap_file, 'r')
d = f[args.statmap_file_subspace_name]
if args.n_loudest is not None:
    n = d['stat'][:].argsort()[::-1][args.n_loudest]
    title = 'Parameters of coincident event ranked %s' % (args.n_loudest + 1)
    caption = 'Parameters of event ranked %s by the search. The figures below show the mini-followup data for this event.' % (args.n_loudest + 1)
elif args.trigger_id is not None:
    n = args.trigger_id
    title = 'Details of coincident trigger'
    caption = 'Parameters of coincident event. The figures below show the mini-followup data for this event.'
else:
    # It shouldn't be possible to get here!
    raise ValueError()

# make a table for the coincident information #################################
headers = ["Coincident ranking statistic",
           "Inclusive IFAR (yr)",
           "Inclusive FAP", 
           "Exclusive IFAR (yr)",
           "Exclusive FAP",
           "Time delay (s)"
          ]

table = numpy.array([['%5.2f' % d['stat'][n], 
                      '%5.2f' % d['ifar'][n], 
                      '%5.2e' % d['fap'][n],
                      '%5.2f' % d['ifar_exc'][n], 
                      '%5.2e' % d['fap_exc'][n],
                      '%5.4f' % (d['time2'][n] - d['time1'][n])
                    ]], dtype=str)

html = pycbc.results.dq.redirect_javascript + \
                                str(pycbc.results.static_table(table, headers))

# make a table for the single detector information ############################
ifo1, ifo2 = f.attrs['detector_1'], f.attrs['detector_2']
idx = {ifo1:d['trigger_id1'][n], ifo2:d['trigger_id2'][n]}

# Store the single detector trigger files keyed by ifo in a dictionary
table = []
files = {}
for fname in args.single_trigger_files:
    f = h5py.File(fname, 'r')
    ifos = f.keys()
    for ifo in ifos:
        files[ifo] = f[ifo]

bank = h5py.File(args.bank_file, 'r')

for ifo in files.keys():
    d = files[ifo]
    i = idx[ifo]
    tid = d['template_id'][i]
    rchisq =  d['chisq'][i] / (d['chisq_dof'][i] * 2 - 2)
    mchirp = (pycbc.pnutils.mass1_mass2_to_mchirp_eta(bank['mass1'][tid], 
                                                      bank['mass2'][tid]))[0]  
                                                      
    time = d['end_time'][:][i]
    utc = lal.GPSToUTC(int(time))[0:6]
                                                                            
    data = [pycbc.results.dq.get_summary_page_link(ifo, utc),
            str(datetime.datetime(*utc)),
            '%.3f'  % time,
            '%5.2f' % d['snr'][i],
            '%5.2f' % pycbc.events.newsnr(d['snr'][i], rchisq),
            '%5.2f' % rchisq,            
            '%i'    % d['chisq_dof'][i],
            '%5.2f' % d['coa_phase'][i],
            '%5.2f' % bank['mass1'][tid],
            '%5.2f' % bank['mass2'][tid],
            '%5.2f' % mchirp,
            '%5.2f' % bank['spin1z'][tid],
            '%5.2f' % bank['spin2z'][tid],
            '%5.2f' % d['template_duration'][i]
           ]
    table.append(data)    

html += str(pycbc.results.static_table(table, pycbc.results.sngl_table_headers))
###############################################################################

pycbc.results.save_fig_with_metadata(html, args.output_file, {},
                        cmd=' '.join(sys.argv),
                        title=title,
                        caption=caption)
