#!/usr/bin/env python
import numpy, h5py, argparse, matplotlib
from matplotlib import colors
matplotlib.use('Agg')
import pylab, pycbc.results

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

# RECALCULE 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', help='HDF format found injection file')
parser.add_argument('--single-injection-file', help='Single detector trigger files from the injection set. One per ifos')
parser.add_argument('--coinc-statistic-file', help='HDF format statistic file')
parser.add_argument('--single-trigger-file', help='Single detector triggers files from the zero lag run')
parser.add_argument('--newsnr-contours', nargs='*', help="List of newsnr values to draw contours at.", default=[])
parser.add_argument('--background-front', action='store_true', default=False)
parser.add_argument('--output-file')
args = parser.parse_args()

import mpld3

# Add the background triggers
f = h5py.File(args.coinc_statistic_file)
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)
ifo = f.keys()[0]
f = f[ifo]
tid = b_tids[ifo][:]
snr = f['snr'][:][tid]

chisq = f['chisq'][:][tid]
l = chisq == 0
chisq_dof = f['chisq_dof'][:][tid]
chisq /= (chisq_dof * 2 - 2) 
chisq[l] = 0.1

fig = pylab.figure()
pylab.scatter(snr, 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)
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)[ifo]
tid = inj_tids[ifo][:]
snr = f['snr'][:][tid]

chisq = f['chisq'][:][tid]
l = chisq == 0
chisq_dof = f['chisq_dof'][:][tid]
chisq /= (chisq_dof * 2 - 2)
chisq[l] = 0.1

pylab.scatter(snr, chisq, c=eff_dist, norm=colors.LogNorm(), 
              marker='^', linewidth=0, label="Injections", 
              zorder= not args.background_front)
    
r = numpy.logspace(numpy.log(chisq.min()*0.9), numpy.log(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')
cb = pylab.colorbar()
cb.set_label('Effective Distance (Mpc)', size='large')
pylab.title('%s Coincident Triggers' % ifo, size='large')
pylab.xlabel('SNR', size='large')
pylab.ylabel('Reduced $\chi^2$', size='large')
pylab.xlim(snr.min()*0.9, snr.max()*1.1)
pylab.ylim(chisq.min()*0.9, chisq.max()*1.1)
pylab.legend(loc='upper left', 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)
pycbc.results.save_fig_with_metadata(fig, args.output_file, 
                                     title='%s background with injections %s' % (ifo.upper(), 'behind' if args.background_front else 'ontop'), 
                                     caption='')
