#!/usr/bin/python3.13 -s
# 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
import h5py
import numpy
import datetime
import argparse
import matplotlib
matplotlib.use('Agg')

import lal

import pycbc.events, pycbc.results, pycbc.pnutils
from pycbc.results import followup
from pycbc.events import stat as pystat
from pycbc.io import hdf
from pycbc import init_logging, add_common_pycbc_options


parser = argparse.ArgumentParser()
add_common_pycbc_options(parser)
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('--title',
    help="Title for the produced snippet. If not given, will default to "
         "pre-sets according to the options given")
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.")
pystat.insert_statistic_option_group(parser,
    default_ranking_statistic='single_ranking_only')


args = parser.parse_args()

init_logging(args.verbose)

if args.trigger_id is not None:
    # Make a mask which is just the trigger of interest
    mask = numpy.array([args.trigger_id])
else:
    mask = None

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

rank_method = pystat.get_statistic_from_opts(args, [args.instrument])

if args.trigger_id is not None:
    # Mask already applied
    pass
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
    )
else:
    raise ValueError("Must give --n-loudest or --trigger-id.")

sds = rank_method.single(sngl_file.trig_dict())
stat = rank_method.rank_stat_single((sngl_file.ifo, sds))
if args.n_loudest is not None:
    # Restrict to only the nth loudest, instead of all triggers up to nth
    # loudest
    l = stat.argsort()
    stat = stat[l[0]]
    sngl_file.apply_mask(l[0])

# 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
    stat_name_long = sngl_stat_name
else:
    # Name would be too long - just call it ranking statistic
    stat_name = 'Ranking Statistic'
    stat_name_long = ' with '.join(
        [args.ranking_statistic, args.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))
###############################################################################

# Set up default titles and the captions for the file
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_long)
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 = args.title if args.title else title,
                        caption = caption)
