#!/usr/bin/python3.11
# 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.
""" Plot the single detector trigger timeseries """
import argparse, logging, pycbc.version, pycbc.results, sys
from pycbc.types import MultiDetOptionAction
from pycbc.events import ranking
from pycbc.io import HFile, SingleDetTriggers
import h5py
import matplotlib; matplotlib.use('Agg'); import pylab
import numpy

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='+',
    action=MultiDetOptionAction, metavar="IFO:FILE",
    help="The HDF format single detector merged trigger files, in "
         "multi-ifo argument format, H1:file1.hdf L1:file2.hdf, etc")
parser.add_argument('--window', type=float, default=10,
    help="Time in seconds around the coincident trigger to plot")
parser.add_argument('--times', nargs='+', type=float,
    action=MultiDetOptionAction, metavar="IFO:GPS_TIME",
    help="The gps times to plot around in multi-ifo argument format, "
         "H1:132341323 L1:132423422")
parser.add_argument('--special-trigger-ids', nargs='+', type=int,
    action=MultiDetOptionAction, metavar="IFO:GPS_TIME",
    help="The set of special trigger ids to plot a star at")
parser.add_argument('--plot-type',
    choices=ranking.sngls_ranking_function_dict, default='snr',
    help="Which single-detector ranking statistic to plot.")
parser.add_argument('--output-file')
parser.add_argument('--log-y-axis', action='store_true')

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

any_data = False

fig = pylab.figure()

min_rank = numpy.inf
for ifo in args.single_trigger_files.keys():
    logging.info("Getting %s triggers", ifo)
    t = args.times[ifo]

    # Identify trigger idxs within window of trigger time
    with HFile(args.single_trigger_files[ifo], 'r') as data:
        idx, _ = data.select(lambda endtime: abs(endtime - t) < args.window,
                             f'{ifo}/end_time', return_indices=True)
        data_mask = numpy.zeros(len(data[f'{ifo}/snr']), dtype=bool)
        data_mask[idx] = True

    if not len(idx):
        # No triggers in this window, add to the legend and continue
        # Make sure it isnt on the plot
        pylab.scatter(-2 * args.window, 0,
                      color=pycbc.results.ifo_color(ifo),
                      marker='x',
                      label=ifo)
        continue

    # Not using bank file, or veto file, so lots of "None"s here.
    trigs = SingleDetTriggers(
        args.single_trigger_files[ifo],
        None,
        None,
        None,
        None,
        ifo,
        premask=data_mask
    )

    any_data = True
    logging.info("Keeping %d triggers in the window", len(idx))

    logging.info("Getting %s", args.plot_type)
    rank = ranking.get_sngls_ranking_from_trigs(trigs, args.plot_type)

    if all(rank == 1):
        # This is the default value to say "downranked beyond consideration"
        # so skip these events
        pylab.scatter(-2 * args.window, 0,
                      color=pycbc.results.ifo_color(ifo),
                      marker='x',
                      label=ifo)
        continue

    pylab.scatter(trigs['end_time'] - t, rank,
                  color=pycbc.results.ifo_color(ifo), marker='x',
                  label=ifo)

    # Minimum rank which hasn't been set to the default of 1
    min_rank = min(min_rank, rank[rank > 1].min())

    if args.special_trigger_ids:
        special_idx = args.special_trigger_ids[ifo]
        if special_idx not in idx:
            logging.info("IDX %d not in kept list",
                         args.special_trigger_ids[ifo])
            continue
        special_red_idx = numpy.where(idx == special_idx)[0]

        pylab.scatter(trigs.trigs_f[f'{ifo}/end_time'][special_idx] - t,
                      rank[special_red_idx], marker='*', s=50, color='yellow')

if args.log_y_axis and any_data:
    pylab.yscale('log')

if not numpy.isinf(min_rank):
    pylab.ylim(ymin=min_rank)

pylab.xlabel('time (s)')
pylab.ylabel(args.plot_type)

pylab.xlim(xmin=-args.window, xmax=args.window)
pylab.legend()
pylab.grid()

logging.info("Saving figure")
pycbc.results.save_fig_with_metadata(fig, args.output_file,
            cmd = ' '.join(sys.argv),
            title = 'Single Detector Trigger Timeseries (%s)' % args.plot_type,
            caption = 'Time series showing the single detector triggers'
                      ' centered around the time of the trigger of interest.',
         )
logging.info("Done!")
