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

""" Plot the single detector trigger timeseries """
import argparse
import logging
import sys
import h5py
import matplotlib
matplotlib.use('Agg')
import pylab
import numpy

from pycbc import init_logging, add_common_pycbc_options
import pycbc.results
from pycbc.types import MultiDetOptionAction
from pycbc.events import ranking
from pycbc.io import HFile, SingleDetTriggers

parser = argparse.ArgumentParser()
add_common_pycbc_options(parser)
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()
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,
            'end_time',
            group=ifo,
            return_data=False,
        )
        data_mask = numpy.zeros(data[ifo]['snr'].size, 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

    trigs = SingleDetTriggers(
        args.single_trigger_files[ifo],
        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)

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

    min_rank = min(min_rank, rank.min())

    if args.special_trigger_ids:
        special_idx = args.special_trigger_ids[ifo]
        if special_idx == None:  # No special trigger for this ifo
            continue
        elif 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. '
                      'Triggers with ranking 1 have been downweighted beyond '
                      'consideration, but may still form insignificant '
                      'events.',
         )
logging.info("Done!")
