#!/usr/bin/env python
#
# Copyright (C) 2017,2018  Heather Fong and Chad Hanna
#
# 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 2 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 h5py
import argparse
import numpy
from glue.ligolw import ligolw
from glue.ligolw import lsctables, param as ligolw_param, array as ligolw_array
from glue.ligolw import utils as ligolw_utils
from glue.ligolw.utils import process as ligolw_process
import lal.series
from lal import rate

@ligolw_array.use_in
@ligolw_param.use_in
@lsctables.use_in
class LIGOLWContentHandler(ligolw.LIGOLWContentHandler):
	pass

def chirpmass(m1, m2):
	return (m1 * m2)**.6 / (m1 + m2)**.2

def schechter_norm(lower, upper, maxM, alpha):
	norm_masses = numpy.linspace(lower, upper, 10000)
	return numpy.sum(schechter(norm_masses, maxM, alpha)) * (norm_masses[1] - norm_masses[0])

def schechter(mass, maxM, alpha):
	return (mass / maxM)**alpha * numpy.exp(-mass / maxM) / maxM

parser = argparse.ArgumentParser(description = "Create analytic mass models for prior weighting of templates")
parser.add_argument("--template-bank", metavar='name', type=str, help='The input template bank file name.', required = True)
parser.add_argument("--output", metavar='name', type=str, help='The output file name', default = "inspiral_mass_model.h5")
parser.add_argument("--model", metavar='name', type=str, help='Mass model. Options are: ligo. If you want another one, submit a patch.')
parser.add_argument("--verbose", help='Be verbose', action="store_true")
options = parser.parse_args()

# Read the template bank file
xmldoc = ligolw_utils.load_filename(options.template_bank, verbose = options.verbose, contenthandler = LIGOLWContentHandler)
sngl_inspiral_table = lsctables.SnglInspiralTable.get_table(xmldoc)

mchirps_of_tmps = chirpmass(sngl_inspiral_table.get_column("mass1"), sngl_inspiral_table.get_column("mass2"))
num_templates = len(mchirps_of_tmps)
num_bins = num_templates / 500
massBA = rate.BinnedDensity(rate.NDBins((rate.LogarithmicBins(min(mchirps_of_tmps)-1e-6, max(mchirps_of_tmps)+1e-6, num_bins),)))
for m in mchirps_of_tmps:
	massBA.count[(m,)] += 1
rate.filter_array(massBA.array, rate.gaussian_window(num_bins / 50., sigma = 5))

# Assign the proper mass probabilities
ids = {}
tmplt_ids = []
mchirps = []
probs = []
bbhnorm = schechter_norm(5., 45, 45., -1.6)
othernorm = schechter_norm(1, 200, 200., -2.35)

for row in sngl_inspiral_table:
	assert row.template_id not in ids
        tmplt_ids.append(int(row.template_id))
	if options.model == "ligo":

		#
		# BNS portion
		#

		# assume a 0.15 solar mass std deviation, this should capture both population distribution and snr effects
		mchirp = chirpmass(row.mass1, row.mass2)
		sigma = 0.15
		mean = 1.2
		bnsprob = 1. / (2 * numpy.pi * sigma**2)**.5 * numpy.exp(-(mchirp - mean)**2 / 2. / sigma**2)

		#
		# BBH portion
		#

		# normalised over 5 -- 45 Msun
		# From: https://www.lsc-group.phys.uwm.edu/ligovirgo/cbcnote/RatesAndSignificance/O1O2CatalogRates
		bbhprob = schechter(mchirp, 45., -1.6) / bbhnorm


		#
		# Other: made up
		#

		otherprob = schechter(mchirp, 200., -2.35) / othernorm

		#
		# Combined
		#

		bns_rate = 916.
		bbh_rate = 56.
		other_rate = 20. # made up
		# FIXME if the noise is ever normalized over mass then we would need the following, but it isn't
		#ids[row.template_id] = (bns_rate * bnsprob + bbh_rate * bbhprob + other_rate * otherprob) / massBA[(mchirp,)] / (bns_rate + bbh_rate + other_rate)
		ids[row.template_id] = (bns_rate * bnsprob + bbh_rate * bbhprob + other_rate * otherprob) / (bns_rate + bbh_rate + other_rate)
		mchirps.append(mchirp)
		probs.append(ids[row.template_id])
	else:
		raise ValueError("Invalid mass model")

coefficients = numpy.zeros((1, 1, max(ids)+1), dtype=float)
for tid in ids:
	coefficients[0,0,tid] = numpy.log(ids[tid] / len(sngl_inspiral_table))

#print coefficients.max()
#import matplotlib
#matplotlib.use('agg')
#from matplotlib import pyplot
#pyplot.loglog(mchirps, probs, "*")
#pyplot.savefig("blah.png")
#

# Write it out
f = h5py.File(options.output, "w")
# put in a dummy interval for the piecewise polynomials in SNR
f.create_dataset("SNR", data = numpy.array([0., 100.]))
f.create_dataset("coefficients", data = coefficients, compression="gzip")
f.create_dataset("event_id", data = numpy.array(tmplt_ids).astype(int))
f.close()