#!/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
from pycbc import pnutils
import pycbc.version

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("--version", action="version",
                    version=pycbc.version.git_verbose_msg)
parser.add_argument('--found-injection-file', required=True,
                    help='HDF format found injection file. Required')
parser.add_argument('--single-injection-file', required=True,
                    help='Single detector trigger files from the injection set. One per ifo')
parser.add_argument('--coinc-statistic-file', required=True,
                    help='HDF format statistic file. Required')
parser.add_argument('--single-trigger-file', required=True,
                    help='Single detector trigger files from the zero lag run.'
                         ' Required')
parser.add_argument('--newsnr-contours', nargs='*', default=[],
                    help="List of newsnr values to draw contours. Optional")
parser.add_argument('--background-front', action='store_true', default=False,
                    help='If set, plot background on top of injections rather '
                         'than vice versa')
parser.add_argument('--colorbar-choice', choices=('effective_spin', 'mchirp',
                    'eta', 'effective_distance', 'mtotal', 'optimal_snr',
                    'redshift'), default='effective_distance',
                    help='Parameter to use for the colorbar. '
                         'Default=effective_distance')
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_exc/trigger_id1']
b_tids[f.attrs['detector_2']] = f['background_exc/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]
m1, m2 = f['injections/mass1'][:][inj_idx], f['injections/mass2'][:][inj_idx]
s1, s2 = f['injections/spin1z'][:][inj_idx], f['injections/spin2z'][:][inj_idx]
mchirp, eta = pnutils.mass1_mass2_to_mchirp_eta(m1, m2)
weighted_spin = pnutils.phenomb_chi(m1, m2, s1, s2)
redshift = f['injections/redshift'][:][inj_idx] if args.colorbar_choice == \
                                                           'redshift' else None

#choices to color the found injections
coloring = {'effective_distance': (eff_dist, "Effective Distance (Mpc)",
                                                             colors.LogNorm()),
            'mchirp': (mchirp, "Chirp Mass", colors.LogNorm()),
            'eta': (eta, "Symmetric Mass Ratio", colors.LogNorm()),
            'effective_spin': (weighted_spin, "Weighted Aligned Spin", None),
            'mtotal': (m1 + m2, "Total Mass", colors.LogNorm()),
            'redshift': (redshift, "Redshift", None)
           }

if 'optimal_snr_1' in f['injections']:
    opt_snr_map = {f.attrs['detector_1']: 'injections/optimal_snr_1',
                   f.attrs['detector_2']: 'injections/optimal_snr_2'}
    opt_snr = f[opt_snr_map[ifo]][:][inj_idx]
    coloring['optimal_snr'] = (opt_snr, 'Optimal SNR', colors.LogNorm())

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=coloring[args.colorbar_choice][0],
              norm=coloring[args.colorbar_choice][2],
              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)

if args.newsnr_contours:
    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(coloring[args.colorbar_choice][1], 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 %s of the injection. Dashed lines show
contours of constant NewSNR.""" % (args.chisq_choice, coloring[args.colorbar_choice][1])
pycbc.results.save_fig_with_metadata(fig, args.output_file, title=title, 
                                     caption=caption, cmd=' '.join(sys.argv),
                                     fig_kwds={'dpi':200})