#!/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.
""" 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
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=['snr', 'newsnr'], default='snr',
    help="Which plot to make; an 'snr' or a newsnr' 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

# organize the single detector triggers files by ifo in a dict
fig = pylab.figure()
for ifo in args.single_trigger_files.keys():
    t = args.times[ifo]

    if args.special_trigger_ids:
        id_loud = args.special_trigger_ids[ifo]

    data = HFile(args.single_trigger_files[ifo], 'r')

    # Identify trigger indices within window
    select_func = lambda *inputs: ((inputs[0] < (t + args.window)) & 
                                   (inputs[0] > (t - args.window)))
    times, snr, chisq, chisq_dof = data.select(select_func, ifo + '/end_time',
                                               ifo + '/snr', ifo + '/chisq',
                                               ifo + '/chisq_dof')
                                 
    # center times on the trigger/chosen time
    times = times - t

    if args.plot_type == 'snr':
        val = snr

        if args.special_trigger_ids:
            top = data[ifo]['snr'][id_loud]

    if args.plot_type == 'newsnr':
        rchisq = chisq / (2 * chisq_dof - 2)
        if len(rchisq) > 0:
            val = ranking.newsnr(snr, rchisq)
        else:
            val = numpy.array([])

        # Get the newnsr of the loudest trigger so we can plot a star there
        if args.special_trigger_ids:
            rchisq = data[ifo]['chisq'][id_loud] / \
                (data[ifo]['chisq_dof'][id_loud] * 2 - 2)
            top = ranking.newsnr(data[ifo]['snr'][id_loud], rchisq)

    if type(val) in [numpy.float32, numpy.float64] or len(val) > 0:
        any_data = True

    pylab.scatter(times, val, color=pycbc.results.ifo_color(ifo), marker='x',
                  label=ifo)

    if args.special_trigger_ids:
        any_data = True
        pylab.scatter([0], [top], marker='*', s=50, color='yellow')

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

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

try:
    if len(val) > 0:
        pylab.ylim(ymin=val.min())
except TypeError:
    pylab.ylim(ymin=val)

pylab.xlim(xmin=-args.window, xmax=args.window)
pylab.legend()
pylab.grid()
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.',
         )
