#!/bin/env python
""" plot a an hdf bank file based on background binning
"""
import argparse, h5py, pycbc.events, pycbc.version, pycbc.results, sys
import numpy
import matplotlib; matplotlib.use('Agg')
import pylab
from itertools import cycle

parser = argparse.ArgumentParser()
parser.add_argument('--version', action='version', version=pycbc.version.git_verbose_msg)
parser.add_argument('--bank-file', help='hdf format template bank file')
parser.add_argument('--background-bins', nargs='+', help='list of background bin format strings')
parser.add_argument('--f-lower',
                    help="Lower frequency cutoff for evaluating template duration. Should"
                         " be equal to the lower cutoff used in inspiral jobs")
parser.add_argument('--output-file', help='output file')
args = parser.parse_args()

fig = pylab.figure()
f = h5py.File(args.bank_file)
data = {'mass1':f['mass1'][:], 'mass2':f['mass2'][:],
        'spin1z':f['spin1z'][:], 'spin2z':f['spin2z'][:]}
if args.f_lower:
    data['f_lower'] = float(args.f_lower)

if args.background_bins:
    locs_dict = pycbc.events.background_bin_from_string(args.background_bins, data)
else:
    locs_dict = {'Template Bank':numpy.arange(0, len(data['mass1']), 1)}

color = cycle(['red', 'green', 'blue', 'purple'])

pylab.grid()
for name in locs_dict:
    locs = locs_dict[name]
    pylab.scatter(data['mass1'][locs], data['mass2'][locs], 
    label=name, linewidth=0, s=1, c=color.next())
    
pylab.legend(loc='upper left', markerscale=5)
pylab.xlabel('Mass1 $M_\odot$')
pylab.ylabel('Mass2 $M_\odot$')
pylab.xlim(data['mass1'].min(), data['mass1'].max())
pylab.ylim(data['mass2'].min(), data['mass2'].max())

pycbc.results.save_fig_with_metadata(fig, args.output_file,
    title="Template Bank and Bins Used to Compute Background",
    caption="This plot shows the template bank in the m1-m2 plane. Each template is colored by the bin that it is placed in to compute the search background. Note that the bins may be chosen in a space higher than two dimensions for spinning templates, causing apparent overlap in the m1-m2 plane.",
    cmd=' '.join(sys.argv))
