#!/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 h5py, argparse, logging, pycbc.version, pycbc.events, pycbc.results, sys
from pycbc.types import MultiDetOptionAction
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:fil2.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 = h5py.File(args.single_trigger_files[ifo], 'r')[ifo]
    times = data['end_time'][:]
    idx = numpy.logical_and(times < t + args.window, times > t - args.window)
                                              
    # center times on the trigger/chosen time
    times = times[idx] - t
    
    if args.plot_type == 'snr':
        val = data['snr'][idx]

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

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

        # Get the newnsr of the loudest trigger so we can plot a star there
        if args.special_trigger_ids:
            rchisq = data['chisq'][id_loud] / (data['chisq_dof'][id_loud] * 2 - 2)
            top = pycbc.events.newsnr(data['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.',           
         )
