#!/usr/bin/env python3
# plot single pastro_model with/without KDE filter

import sys
from gstlal import rio;
import numpy
from ligo.lw import ligolw
from ligo.lw import utils as ligolw_utils
from ligo.lw.utils import process as ligolw_process
from lal import rate
from optparse import OptionParser
from gstlal import far
import pastro

# plotting modules
import matplotlib
from matplotlib import figure
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from mpl_toolkits.axes_grid1 import make_axes_locatable


def parse_command_line():
  parser = OptionParser()
  parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
  parser.add_option("--output", help = "output_directory")
  parser.add_option("--category", help = "category")
  parser.add_option("--mc-kde", help = "kde bins for mc")
  parser.add_option("--lr-kde", help = "kde bins for lr")
  parser.add_option("--sigma-kde", help = "sigma for gaussian distribution")
  parser.add_option("--apply-KDE", help = "0 is False, 1 is True if apply KDE")
  options, filenames = parser.parse_args()
  process_params = dict(options.__dict__)

  return options, process_params, filenames

def plot_pdf(p_x_c, category, KDE=True):

    fontsize = 20
    fig = figure.Figure()
    FigureCanvas(fig)
    fig.set_size_inches((15, 12))

    axes = fig.gca()
    axes.set_xlabel('mc bin index', fontsize=fontsize)
    axes.set_ylabel('lr bin index', fontsize=fontsize)
    # axes.set_xticks()
    axes.tick_params(axis='both', which='major', labelsize=15)
    divider= make_axes_locatable(axes)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    
    if KDE:
      fn = category+' (mc, lr, sig)=('+str(options.mc_kde)+','+str(options.lr_kde)+','+str(options.sigma_kde)+')'
    else:
      fn = category + '_no_KDE'
    print(fn)
    axes.set_title(fn, fontsize=fontsize)
    
    im = axes.imshow(p_x_c)

    fig.colorbar(im, cax=cax)
    return fig

options, process_params, filenames = parse_command_line()
model = pastro.pastro_model.from_h5(filenames[0])

categories = ['Terrestrial', 'BNS', 'NSBH', 'BBH']
mc_KDE_bins = {}
lr_KDE_bins = {}
sigma_KDE = {}

for category in categories:
  mc_KDE_bins[category] = float(options.mc_kde)
  lr_KDE_bins[category] = float(options.lr_kde)
  sigma_KDE[category] = float(options.sigma_kde)

if int(options.apply_KDE) == 1:
  KDE = True
elif int(options.apply_KDE) == 0:
  KDE = False
else:
  raise ValueError('wrong index for KDE')

kde_pdf = model.finalize_model(mc_KDE_bins, lr_KDE_bins, sigma_KDE, KDE=KDE)
fig = plot_pdf(kde_pdf[options.category], options.category, KDE=KDE)

if KDE:
  fig.savefig(options.output+options.category+'_'+str(options.mc_kde)+'_'+str(options.lr_kde)+'_'+str(options.sigma_kde)+'.pdf')
else:
  fig.savefig(options.output+options.category+'_no_kde.pdf')
