#!/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 sngl event"""
import sys, h5py, matplotlib, numpy, datetime, argparse
matplotlib.use('Agg')
import lal
import pycbc.version, pycbc.events, pycbc.results, pycbc.pnutils
from pycbc.results import followup
from pycbc.events import stat
from pycbc.io import hdf


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-file', required=True,
    help="HDF format single detector trigger files for the full "
         "data run")
parser.add_argument('--bank-file', required=True,
    help="HDF format template bank file")
parser.add_argument('--output-file')
# CURRENTLY UNUSED, but may be needed if using foreground censor
#parser.add_argument('--statmap-file',
#    help="The HDF format clustered coincident statmap file containing the "
#         "result triggers.")
parser.add_argument('--veto-file',
    help="The veto file to be used if vetoing triggers.  Optional")
parser.add_argument('--veto-segment-name',
    help="If using veto file, the name of the segments to use as a veto.")
parser.add_argument("--instrument", help="Name of ifo (e.g. H1)")
trig_args = parser.add_mutually_exclusive_group(required=True)
trig_args.add_argument("--n-loudest", type=int, default=None,
    help="Examine the n'th loudest trigger (loudest is n=0). Must supply "
         "either this or --trigger-id.")
trig_args.add_argument("--trigger-id", type=int, default=None,
    help="Use the trigger with this ID in the HDF file. Must supply either "
         "this option or --n-loudest.")
parser.add_argument('--include-summary-page-link', action='store_true',
    help="If given, will include a link to the DQ summary page")
parser.add_argument('--include-gracedb-link', action='store_true',
    help="If given, will provide a link to search GraceDB for events "
         "within a 3s window around the coincidence time.")
parser.add_argument('--significance-file',
    help="If given, will search for this trigger's id in the file to see if "
         "stat and p_astro values exists for this trigger.")
stat.insert_statistic_option_group(parser)


args = parser.parse_args()

pycbc.init_logging(args.verbose)

# Get the single-ifo triggers
sngl_file = hdf.SingleDetTriggers(args.single_trigger_file, args.bank_file,
                args.veto_file, args.veto_segment_name, None, args.instrument)

rank_method = stat.get_statistic_from_opts(args, [args.instrument])
if args.trigger_id is not None:
    sngl_file.mask = [args.trigger_id]
    trig_id = args.trigger_id
    sngl_file.mask_to_n_loudest_clustered_events(rank_method, n_loudest=1)
    stat = sngl_file.stat[0]
elif args.n_loudest is not None:
    # Cluster by a ranking statistic and retain only the loudest n clustered
    # triggers
    sngl_file.mask_to_n_loudest_clustered_events(rank_method,
                                                 n_loudest=args.n_loudest+1)
    # Restrict to only the nth loudest, instead of all triggers up to nth
    # loudest
    stat = sngl_file.stat
    l = stat.argsort()
    stat = stat[l[0]]
    sngl_file.mask = sngl_file.mask[l[0]]
else:
    raise ValueError("Must give --n-loudest or --trigger-id.")

# make a table for the single detector information ############################
time = sngl_file.end_time
utc = lal.GPSToUTC(int(time))[0:6]

# Headers here will contain the list of headers that will appear in the
# resulting single-detector summary table.
headers = []
# Data will store the list of values that will appear in the resulting
# single-detector summary table. There will only ever be one row in this table
# that corresponding to the selected trigger. So this will be a list of a
# single list that will hold the values to go into the table.
data = [[]]

# DQ summary link
if args.include_summary_page_link:
    data[0].append(pycbc.results.dq.get_summary_page_link(args.instrument, utc))
    headers.append("Detector&nbsp;status")

# End times
data[0].append(str(datetime.datetime(*utc)))
data[0].append('%.3f' % time)
headers.append("UTC")
headers.append("End&nbsp;time")

# SNR and statistic
data[0].append('%5.2f' % sngl_file.snr)
data[0].append('%5.2f' % sngl_file.get_column('coa_phase'))
data[0].append('%5.2f' % stat)
headers.append("&rho;")
headers.append("Phase")
# Determine statistic naming
if args.sngl_ranking == "newsnr":
    sngl_stat_name = "Reweighted SNR"
elif args.sngl_ranking == "newsnr_sgveto":
    sngl_stat_name = "Reweighted SNR (+sgveto)"
elif args.sngl_ranking == "newsnr_sgveto_psdvar":
    sngl_stat_name = "Reweighted SNR (+sgveto+psdvar)"
elif args.sngl_ranking == "snr":
    sngl_stat_name = "SNR"
else:
    sngl_stat_name = args.sngl_ranking

if args.ranking_statistic in ["quadsum", "single_ranking_only"]:
    stat_name = sngl_stat_name
else:
    stat_name = '_with_'.join(ranking_statistic, sngl_ranking)

headers.append(stat_name)

# Signal-glitch discrimators
data[0].append('%5.2f' % sngl_file.rchisq)
data[0].append('%i' % sngl_file.get_column('chisq_dof'))
headers.append("&chi;<sup>2</sup><sub>r</sub>")
headers.append("&chi;<sup>2</sup>&nbsp;bins")
try:
    data[0].append('%5.2f' % sngl_file.sgchisq)
    headers.append("sg&chi;<sup>2</sup>")
except:
    pass
try:
    data[0].append('%5.2f' % sngl_file.psd_var_val)
    headers.append("PSD var")
except:
    pass

# Template parameters
data[0].append('%5.2f' % sngl_file.mass1)
data[0].append('%5.2f' % sngl_file.mass2)
data[0].append('%5.2f' % sngl_file.mchirp)
data[0].append('%5.2f' % sngl_file.spin1z)
data[0].append('%5.2f' % sngl_file.spin2z)
data[0].append('%5.2f' % sngl_file.template_duration)
headers.append("m<sub>1</sub>")
headers.append("m<sub>2</sub>")
headers.append("M<sub>c</sub>")
headers.append("s<sub>1z</sub>")
headers.append("s<sub>2z</sub>")
headers.append("Duration")

if args.significance_file and not args.n_loudest:
    with hdf.HFile(args.significance_file, 'r') as sig_f:
        trigger_ids = sig_f['foreground'][args.instrument]['trigger_id'][:]
        sngl_stat = sig_f['foreground/stat'][:]
        sngl_pastro = sig_f['foreground/p_astro_exc'][:]

    if trig_id in trigger_ids:
        trig_idx = numpy.nonzero(trigger_ids == trig_id)[0][0]
        headers.append("Single-detector statistic")
        data[0].append("%.2f" % sngl_stat[trig_idx])
        headers.append("Single-detector p_astro")
        pastro_form = "%.3f" if sngl_pastro[trig_idx] >= 0.001 else "%.3e"
        data[0].append(pastro_form % sngl_pastro[trig_idx])

# GraceDB search link
if args.include_gracedb_link:
    gdb_search_link = followup.get_gracedb_search_link(time)
    headers.append("GraceDB Search Link")
    data[0].append(gdb_search_link)

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

if args.n_loudest:
    title = 'Parameters of single-detector event ranked %s' \
        % (args.n_loudest + 1)
    caption = 'Parameters of the single-detector event ranked number %s by %s. The figures below show the mini-followup data for this event.' % (args.n_loudest + 1, stat_name)
else:
    title = 'Parameters of single-detector event'
    caption = 'Parameters of the single-detector event. The figures below show the mini-followup data for this event.'

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