#!/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 output of pycbc_single_template """

import argparse
import sys
import h5py
import numpy
import matplotlib
matplotlib.use('Agg')
import pylab
import pycbc.results
from pycbc.events import ranking


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--version', action='version',
    version=pycbc.version.git_verbose_msg)
parser.add_argument('--verbose')
parser.add_argument('--single-template-file', required=True,
    help="HDF file containing the SNR and CHISQ timeseries. "
         " The output of pycbc_single_template")
parser.add_argument('--window', type=float, required=True,
    help="The time in seconds to plot")
parser.add_argument('--event-time', type=float,
                    help='GPS time to use as the center of the plot')
parser.add_argument('--output-file', required=True)
parser.add_argument('--log-y', action='store_true',
                    help='Use log scale for y-axis')
parser.add_argument('--plot-title',
                    help="If given, use this as the plot title")
parser.add_argument('--plot-caption',
                    help="If given, use this as the plot caption")

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

# The event is chosen from the input of pycbc_single template, so we
# just extract the parameters here
f = h5py.File(args.single_template_file, 'r')

try:
    delta_t = f['snr'].attrs['delta_t']
    start_time = f['snr'].attrs['start_time']
except:
    pylab.text(0.5, 0.5, 'no triggers found')
    pylab.savefig(args.output_file)
    sys.exit()

if args.event_time is not None:
    time = args.event_time
elif 'event_time' in f.attrs:
    time = f.attrs['event_time']
else:
    raise ValueError('Event time is neither specified nor found in the HDF file')

ifo = f.attrs['ifo']
center = int((time - start_time) / delta_t)

# Determine where in the timeseries we need to plot
left = center - int(args.window / delta_t)
right = center + int(args.window / delta_t)

left = max(left,0)
right = min(right,len(f['snr']))

snr = abs(f['snr'][left:right][:])
chisq = f['chisq'][left:right][:]

rang = (numpy.arange(0, len(snr), 1) - (center - left)) * delta_t
newsnr = ranking.newsnr(snr, chisq)

fig, ax1 = pylab.subplots()
ax1.plot(rang, snr, color='blue', label='SNR')
ax1.plot(rang, newsnr, color='purple', label='NewSNR')
ax1.set_ylabel('SNR')
if args.log_y:
    ax1.set_yscale('log')
ax1.legend(loc="upper left")
ax1.set_xlabel('Time - %.3f (s)' % time)

ax2 = ax1.twinx()
ax2.plot(rang, chisq, color='green', label='$\\chi^2_{\\rm r}$')
ax2.set_ylabel('Reduced $\\chi^2$')
if args.log_y:
    ax2.set_yscale('log')
ax2.legend(loc="upper right")

if args.plot_title is None:
    args.plot_title = '%s: SNR and chi^2 time series' % ifo
if args.plot_caption is None:
    args.plot_caption = ''

if 'command_line' in f.attrs:
    cmd = f.attrs['command_line'].decode() + '\\n\\n' + ' '.join(sys.argv)
else:
    cmd = ' '.join(sys.argv)

pycbc.results.save_fig_with_metadata(fig, args.output_file, 
                             cmd=cmd, fig_kwds={'dpi': 150},
                             title=args.plot_title,
                             caption=args.plot_caption)
