#!/usr/bin/env python

# Copyright (C) 2015 Christopher M. Biwer
#
# 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.

import argparse
import glob
import math
import os
from itertools import groupby

import numpy as np
import matplotlib as mpl; mpl.use('Agg')
import matplotlib.pylab as plt

from glue import lal
from glue import segments

from pylal import SnglInspiralUtils

# parse command line
parser = argparse.ArgumentParser(usage='pycbc_plot_bank [--options]',
                  description="Plot template bank points that have triggers.")
parser.add_argument('--trigger-files', type=str, nargs="+",
                  help='Input xml files with SnglInspiral triggers.')
parser.add_argument('--output-file', type=str,
                  help='Output image file.')
parser.add_argument('--gps-start-time', type=int,
                  help='Time to start accepting triggers for plotting.')
parser.add_argument('--gps-end-time', type=int,
                  help='Time to stop accepting triggers for plotting.')
opts = parser.parse_args()

# read inspiral triggers
inspiralTable = SnglInspiralUtils.ReadSnglInspiralFromFiles(opts.trigger_files)

# put data into dict where key is (mtotal,eta) both rescaled and value is counts
scale_factor = 1000
bank = {}
for trig in inspiralTable:
    trig_time = trig.end_time + trig.end_time_ns * 10**(-9)
    trig_mtotal = int(trig.mtotal * scale_factor)
    trig_eta = int(trig.eta * scale_factor)
    if trig_time >= opts.gps_start_time and trig_time < opts.gps_end_time:
        if (trig_mtotal,trig_eta) in bank.keys():
            bank[(trig_mtotal,trig_eta)] += 1
        else:
            bank[(trig_mtotal,trig_eta)] = 1

# put data into lists for plotting
param_x = []
param_y = []
param_z = []
for key, count in bank.iteritems():
    mtotal, eta = key
    mtotal = float(mtotal) / scale_factor
    eta = float(eta) / scale_factor
    param_x.append(mtotal)
    param_y.append(eta)
    param_z.append(math.log(count, 10))

# turn data into pylab arrays
x = plt.array(param_x)
y = plt.array(param_y)
z = plt.array(param_z)

# create pylab figure
fig  = plt.figure(1, figsize=(8.1,8.1))

# plot data
p = plt.scatter(x, y, c=z)

# format the plot
cb = plt.colorbar()
cb.set_label('log10(count)')
plt.ylim(min(param_y) * 0.9, max(param_y) * 1.1)
plt.ylabel('Eta')
plt.xlim(min(param_x) * 0.9, max(param_x) * 1.1)
plt.xlabel('Total mass (solar masses)')

# save plot
plt.savefig(opts.output_file)
plt.close()