#!/usr/bin/env python

import sys
import numpy, h5py, argparse, matplotlib
from matplotlib import colors
matplotlib.use('Agg')
import pylab, pycbc.results
from pycbc.io import get_chisq_from_file_choice, chisq_choices

def snr_from_chisq(chisq, newsnr, q=6.):
    snr = numpy.zeros(len(chisq)) + float(newsnr)
    ind = numpy.where(chisq > 1.)[0]
    snr[ind] = float(newsnr) / ( 0.5 * (1. + chisq[ind] ** (q/2.)) ) ** (-1./q)
    return snr

# FIXME: RECALCULATE THIS FROM DISTANCE AND ORIENTATION
eff_map = {'H1':'eff_dist_h', 'L1':'eff_dist_l', 'V1':'eff_dist_v'} 

parser = argparse.ArgumentParser()
parser.add_argument('--found-injection-file', required=True,
                    help='HDF format found injection file')
parser.add_argument('--single-injection-file', required=True,
                    help='Single detector trigger files from the injection set. One per ifos')
parser.add_argument('--coinc-statistic-file', required=True,
                    help='HDF format statistic file')
parser.add_argument('--single-trigger-file', required=True,
                    help='Single detector triggers files from the zero lag run')
parser.add_argument('--newsnr-contours', nargs='*', default=[],
                    help="List of newsnr values to draw contours at.")
parser.add_argument('--background-front', action='store_true', default=False)
parser.add_argument('--chisq-choice', choices=chisq_choices,
                    default='traditional',
                    help='Which chisquared to plot. Default=traditional')
parser.add_argument('--output-file', required=True)
args = parser.parse_args()

import mpld3

# Add the background triggers
f = h5py.File(args.coinc_statistic_file, 'r')
b_tids = {}
b_tids[f.attrs['detector_1']] = f['background/trigger_id1']
b_tids[f.attrs['detector_2']] = f['background/trigger_id2']

f = h5py.File(args.single_trigger_file, 'r')
ifo = f.keys()[0]
f = f[ifo]
tid = b_tids[ifo][:]
bkg_snr = f['snr'][:][tid]

bkg_chisq = get_chisq_from_file_choice(f, args.chisq_choice)[tid]

fig = pylab.figure()
pylab.scatter(bkg_snr, bkg_chisq, marker='o', color='black',
              linewidth=0, s=2.0, label='Background', 
              zorder=args.background_front)

# Add the found injection points
f = h5py.File(args.found_injection_file, 'r')
inj_tids = {}
inj_tids[f.attrs['detector_1']] = f['found_after_vetoes/trigger_id1']
inj_tids[f.attrs['detector_2']] = f['found_after_vetoes/trigger_id2']

inj_idx = f['found_after_vetoes/injection_index'][:]
eff_dist = f['injections'][eff_map[ifo]][:][inj_idx]

f = h5py.File(args.single_injection_file, 'r')[ifo]
tid = inj_tids[ifo][:]
if len(tid):
    inj_snr = f['snr'][:][tid]
    inj_chisq = get_chisq_from_file_choice(f, args.chisq_choice)[tid]
else:
    inj_snr = numpy.array([])
    inj_chisq = numpy.array([])

pylab.scatter(inj_snr, inj_chisq, c=eff_dist, norm=colors.LogNorm(),
              marker='^', linewidth=0, label="Injections", 
              zorder= not args.background_front)

try:
    r = numpy.logspace(numpy.log(min(bkg_chisq.min(), inj_chisq.min())*0.9),
                   numpy.log(max(bkg_chisq.max(), inj_chisq.max())*1.1), 200)
except ValueError:
    # Allow code to continue in the absence of injection triggers
    r = numpy.logspace(numpy.log(bkg_chisq.min()*0.9),
                       numpy.log(bkg_chisq.max()*1.1), 200)
for cval in args.newsnr_contours:
    snrv = snr_from_chisq(r, cval)
    pylab.plot(snrv, r, '--', color='grey', linewidth=1)
    
ax = pylab.gca()
ax.set_xscale('log')
ax.set_yscale('log')
try:
    cb = pylab.colorbar()
    cb.set_label('Effective Distance (Mpc)', size='large')
except TypeError:
    # Catch case of no injection triggers
    if len(inj_chisq):
        raise
pylab.title('%s Coincident Triggers' % ifo, size='large')
pylab.xlabel('SNR', size='large')
pylab.ylabel('Reduced $\chi^2$', size='large')
try:
    pylab.xlim(min(inj_snr.min(),bkg_snr.min())*0.9,
               max(inj_snr.max(),bkg_snr.max())*1.4)
    pylab.ylim(min(bkg_chisq.min(), inj_chisq.min())*0.7,
               max(bkg_chisq.max(), inj_chisq.max())*1.4)
except ValueError:
    # Raised if no injection triggers
    pass
pylab.legend(loc='lower right', prop={'size': 12})
pylab.grid(which='major', ls='solid', alpha=0.7, linewidth=.5)
pylab.grid(which='minor', ls='solid', alpha=0.7, linewidth=.1)

title = '%s %s chisq vs SNR. Background with injections %s' \
        % (ifo.upper(), args.chisq_choice,
           'behind' if args.background_front else 'ontop')
caption = """Distribution of SNR and %s chi-squared veto for single detector
triggers. Black points are background triggers. Triangles are injection
triggers colored by the effective distance of the injection. Dashed lines show
contours of constant NewSNR.""" % args.chisq_choice
pycbc.results.save_fig_with_metadata(fig, args.output_file, title=title, 
                                     caption=caption, cmd=' '.join(sys.argv),
                                     fig_kwds={'dpi':300})
