#!/usr/bin/python3.13 -s
""" Make histograms of single detector triggers
"""

import numpy
import argparse
import logging
import sys
from itertools import cycle
from matplotlib import use
use('Agg')
from matplotlib import pyplot

import pycbc
import pycbc.results
import pycbc.io
from pycbc.events import background_bin_from_string, veto, ranking

parser = argparse.ArgumentParser(description=__doc__)
pycbc.add_common_pycbc_options(parser)
parser.add_argument('--trigger-file', required=True,
                    help="Combined single detector hdf trigger file")
parser.add_argument('--veto-file',
                    help="segment xml file, indicates which triggers to ignore")
parser.add_argument('--segment-name',
                    help="name of segment list in the veto file")
parser.add_argument('--x-var',
                    help="name of value to histogram")
parser.add_argument('--output-file')
parser.add_argument('--bank-file', default="",
                    help='template bank hdf file')
parser.add_argument('--background-bins', nargs='+',
                    help='list of background bin format strings')
parser.add_argument('--bins', default=100, type=int,
                    help="number of bins in histogram")
parser.add_argument('--x-max', type=float)
parser.add_argument('--x-min', type=float, default=6,
                    help="Minimum x-value. Default 6")
parser.add_argument('--special-time', type=float,
                    help="plot triggers within +-1s of a given time in a "
                         "different color (black)")
args = parser.parse_args()

# sanity check command line options
if args.special_time and args.background_bins:
    raise ValueError("Cannot use both --special-time and --background-bins")

# setup logging
pycbc.init_logging(args.verbose)

# read trigger file to determine IFO
f = pycbc.io.HFile(args.trigger_file, 'r')
ifo = tuple(f.keys())[0]

# Apply presort if we can; if we can't, load everything but it may be
# inefficient
xv = args.x_var if args.x_var in ranking.sngls_ranking_function_dict else None
xt = args.x_min if args.x_var in ranking.sngls_ranking_function_dict else None

# read single-detector triggers
trigs = pycbc.io.SingleDetTriggers(
    args.trigger_file,
    ifo,
    bank_file=args.bank_file,
    veto_file=args.veto_file,
    segment_name=args.segment_name,
    filter_rank=xv,
    filter_threshold=xt,
)

# get x values and find the maximum x value
val = getattr(trigs, args.x_var)
x_max = args.x_max if args.x_max else val.max() * 1.1

# create a figure to add a histogram
fig = pyplot.figure(0)

# command line says to plot triggers around a certain time
if args.special_time:
    time = getattr(trigs, 'end_time')
    val_special = val[abs(time-args.special_time) < 1]
    val_boring = val[abs(time-args.special_time) >= 1]
    binvals = numpy.linspace(args.x_min, x_max, args.bins, endpoint=True)
    pyplot.hist([val_boring, val_special], bins=binvals, histtype='stepfilled',
               stacked=True, color=[pycbc.results.ifo_color(ifo), 'k'],
          label=['Other times', 'Triggers near %i' % (int(args.special_time))])
    pyplot.legend(loc='upper right')

# command line says to plot triggers in each background bin
elif args.background_bins:

    # get a dict where key is each bin's name and value is a list of indexes
    # for the triggers in that bin
    bank_data = {'mass1' : trigs.mass1,
                 'mass2' : trigs.mass2,
                 'spin1z' : trigs.spin1z,
                 'spin2z' : trigs.spin2z,
    }
    locs_dict = background_bin_from_string(args.background_bins, bank_data)
 
    # get a list of bin names and a corresponding list for x values
    loc_bin_keys = [key for key in locs_dict.keys()]
    loc_bin_vals = [val[locs_dict[key]] for key in loc_bin_keys]

    # assign a color for each bin
    color_cycle = cycle(['red', 'green', 'blue', 'black', 'magenta', 'cyan'])
    loc_bin_colors = [next(color_cycle) for key in loc_bin_keys]

    # get number of overflows for each background bin
    loc_bin_overflows = [len(vals[vals>=x_max]) for vals in loc_bin_vals]
    num_bins = len(loc_bin_overflows)

    # remove overflow triggers from plotting
    loc_bin_vals = [vals[vals<x_max] for vals in loc_bin_vals]

    # plot histograms
    pyplot.hist(loc_bin_vals, bins=args.bins, histtype='step',
                stacked=False, label=loc_bin_keys, color=loc_bin_colors)
    pyplot.legend(loc='upper right')

    # plot overflow bin
    rects = pyplot.bar(num_bins*[x_max+0.05],
                      loc_bin_overflows, .5,
                      facecolor='None', edgecolor=loc_bin_colors)

    # write text for overflow bin
    text_left = rects[0].get_x()
    text_height = 1.10 * max([rect.get_height() for rect in rects])
    pyplot.text(text_left, text_height, '%s+'%x_max)

# plot all triggers in a single label
else:
    pycbc.results.hist_overflow(val, x_max, bins=args.bins,
                                  color=pycbc.results.ifo_color(ifo))

# format plot
ax = pyplot.gca()
ax.set_yscale('log')
pyplot.ylabel('Number of triggers')
pyplot.xlabel(args.x_var)
pyplot.ylim(ymin=.1)

# set x lower limit on plot
if args.x_min:
    pyplot.xlim(xmin=args.x_min)

# set x upper limit on plot
if args.x_max:
    if len(numpy.where(val > args.x_max)[0]):
        overflow = 1
    else:
        overflow = 0

    pyplot.xlim(xmax=args.x_max + overflow)

# add a grid to the plot
pyplot.grid()

# add meta data and save figure
pycbc.results.save_fig_with_metadata(fig, args.output_file,
                title = '%s: %s histogram of single detector triggers' % (ifo, args.x_var),
                caption = 'Histogram of single detector triggers',
                cmd = ' '.join(sys.argv))

