#!/usr/bin/env python
#
# Copyright (C) 2019  Kipp Cannon, 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.
#

#
# =============================================================================
#
#                                   Preamble
#
# =============================================================================
#

import sys
from optparse import OptionParser
import numpy
import scipy
from scipy.stats import logistic

from lal import series
from lal import MSUN_SI

from ligo.lw import ligolw
from ligo.lw import lsctables
from ligo.lw import array as ligolw_array
from ligo.lw import param as ligolw_param
from ligo.lw import utils as ligolw_utils
from ligo.lw.utils import process as ligolw_process
from lalinspiral.sbank.waveforms import SnglInspiralTable
from lalsimulation import SimIMRSEOBNRv4ROMTimeOfFrequency as lalsim_chirptime

from gstlal import far
from gstlal import svd_bank
from gstlal import templates

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

__author__ = "Chad Hanna <chad.hanna@ligo.org>"
__version__ = "git id %s" % ""	# FIXME
__date__ = ""	# FIXME


#
# =============================================================================
#
#                                 Command Line
#
# =============================================================================
#



def chispace(start, stop, numpoints):
	x = numpy.linspace(0., 5.5, numpoints)
	return start + ((logistic.cdf(x) - 0.5) / (logistic.cdf(x[-1]) - 0.5)) * (stop - start)

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

parser = OptionParser(
	version = "Name: %%prog\n%s" % "" # FIXME
)
parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
parser.add_option("--instrument", action = "append", help = "Append to a list of instruments to create dist stats for.  List must be whatever instruments you intend to analyze.")
parser.add_option("--output-name", type = "string", default = "gstlal_inspiral_grid_bank.xml.gz", help = "Specify output file name. Default gstlal_inspiral_grid_bank.xml.gz")
parser.add_option("--psd-xml", type = "string", help = "Specify a PSD to use for computing template bandwidth. Required if df is bandwidth")
parser.add_option("--m1", type = "string", help = "Specify m1 as start:stop:numpoints")
parser.add_option("--m2", type = "string", help = "Specify m2 as start:stop:numpoints")
parser.add_option("--chi", type = "string", help = "Specify chi as start:stop:numpoints")
parser.add_option("--bandwidth-min", type = "float", help = "Specify minimum bandwidth to keep")
parser.add_option("--min-q", type = "float", default = 1.0, help = "Specify min q, default 1.0")
parser.add_option("--max-q", type = "float", default = float("inf"), help = "Specify max q, default inf")
parser.add_option("--min-mtot", type = "float", default = 0., help = "Specify min mtotal, default 0")
parser.add_option("--max-mtot", type = "float", default = float("inf"), help = "Specify max mtotal, default inf")
options, filenames = parser.parse_args()


if not options.instrument:
	raise ValueError("must specify at least one --instrument")
options.instrument = set(options.instrument)

psd = {}
if options.psd_xml:
	for ifo, p in series.read_psd_xmldoc(ligolw_utils.load_filename(options.psd_xml, verbose = options.verbose, contenthandler = series.PSDContentHandler)).items():
		f = numpy.arange(len(p.data.data)) * p.deltaF
		psd[ifo] = scipy.interpolate.interp1d(f, p.data.data)

m1start, m1stop, m1num = [float(x) for x in options.m1.split(":")]
m2start, m2stop, m2num = [float(x) for x in options.m2.split(":")]
chistart, chistop, s1num = [float(x) for x in options.chi.split(":")]

numtmps = 0
m1s = []
m2s = []
s1s = []
s2s = []
durs = []
for m1 in numpy.logspace(numpy.log10(m1start), numpy.log10(m1stop), m1num):
	for m2 in numpy.logspace(numpy.log10(m2start), numpy.log10(m2stop), m2num):
		for chi in chispace(chistart, chistop, s1num):
			s1 = chi
			s2 = (chi * (m1 + m2) - m1 * s1) / m2
			if m1 < m2:
				continue
			if m1 / m2 > options.max_q or m1 / m2 < options.min_q:
				continue
			if m1 + m2 > options.max_mtot or m1 + m2 < options.min_mtot:
				continue
			if options.bandwidth_min is not None:
				bw = templates.bandwidth(m1, m2, s1, s2, f_min = 10.0, f_max = 1024., delta_f = 0.25, psd = psd[ifo])
				if bw < options.bandwidth_min:
					continue
			dur = lalsim_chirptime(10.0, MSUN_SI*m1, MSUN_SI*m2, s1, s2)
			numtmps += 1
			m1s.append(m1)
			m2s.append(m2)
			s1s.append(s1)
			s2s.append(s2)
			durs.append(dur)
			print m1, m2, s1, s2, numtmps

# prepare a new XML document for writing template bank
xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())
#sngl_inspiral_columns = ("process:process_id", "mass1", "mass2", "spin1x", "spin1y", "spin1z", "spin2x", "spin2y", "spin2z", "mchirp")
lsctables.SnglInspiralTable.RowType = SnglInspiralTable
tbl = lsctables.New(lsctables.SnglInspiralTable)
xmldoc.childNodes[-1].appendChild(tbl)
# FIXME make a real process table
process = ligolw_process.register_to_xmldoc(xmldoc, sys.argv[0], {})
ligolw_process.set_process_end_time(process)

if options.verbose:
	print >> sys.stderr, "Writing output document"

for n, (m1, m2, s1, s2, dur) in enumerate(zip(m1s, m2s, s1s, s2s, durs)):
	row = SnglInspiralTable()
	row.mass1, row.mass2, row.spin1x, row.spin1y, row.spin1z, row.spin2x, row.spin2y, row.spin2z, row.mchirp, row.template_duration = (m1, m2, 0., 0., s1, 0., 0., s2, mchirp(m1, m2), dur)
	row.event_id = n
	row.ifo = "H1" # FIXME
	row.process_id = process.process_id
	tbl.append(row)

ligolw_utils.write_filename(xmldoc, options.output_name,  gz=options.output_name.endswith("gz"))

import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot
pyplot.loglog(m1s, m2s, '*')
pyplot.savefig("bw.png")