#!/usr/bin/env python
#
# Copyright (C) 2012  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 math
import bisect
import cgi
import cgitb
import os
os.environ["MPLCONFIGDIR"] = "/tmp"
import sys
import matplotlib
matplotlib.use('Agg')
from matplotlib import figure
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import matplotlib.pyplot as plt
import numpy
cgitb.enable()
form = cgi.FieldStorage()
import lal
from glue.ligolw import ligolw
from glue.ligolw import array as ligolw_array
from glue.ligolw import lsctables
from glue.ligolw import param as ligolw_param
from glue.ligolw import utils as ligolw_utils
from pylal import series as lalseries
from gstlal import far
from gstlal import reference_psd
import copy
import StringIO
import base64
import urlparse

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

## @file gstlal_llcbcnode
# This program will monitor the output of a single job in the gstlal inspiral
# low latency analysis; See gstlal_llcbcnode for help and usage.

## @package gstlal_llcbcnode
#
# This program is designed to be placed in the cgi-bin directory of the user's
# public_html directory on the cluster that is running the gstlal inspiral low
# latency analysis
#
# ## USAGE:
# This program is never meant to be executed by a user, but rather on a
# webserver via a url such as:
# 
#               https://ldas-jobs.ligo.caltech.edu/~gstlalcbc/cgi-bin/gstlal_llcbcnode?dir=/path/to/analysis/dir/&id=<jobid>&url=/path/to/likelihood.xml
#
# e.g.,
#
#               https://ldas-jobs.ligo.caltech.edu/~gstlalcbc/cgi-bin/gstlal_llcbcnode?dir=/mnt/qfs3/gstlalcbc/engineering/5/bns_trigs_40Hz&id=0009&url=/mnt/qfs3/gstlalcbc/engineering/5/bns_trigs_40Hz/0009_likelihood.xml
#
#
# ## Interpretation of the output page
#
# ### Trigger stats
#
# \image html gstlal_llcbcnode01.png
#
# - The leftmost plot is trigger latency vs time.  It should be flat. If it is
# rising as a function of time that indicates that the job may be falling
# behind.  Occasional spikes are expected 
#
# - The next plot is a trigger latency histogram over the entire duration of
# the run.
#
# - The next plot is a SNR vs time. It should be relatively flat and below 10
# for well behaved noise.  An occasional spike might be caused from a signal or
# a glitch, but long term features indicate poor data quality
#
# - The rightmost plot is the maximum RAM usage history. It can only ever go
# up, but it is a problem if it nears the maximum RAM available on the machine.
#
# ### Live time
#
# \image html gstlal_llcbcnode02.png
#
# These per IFO pie charts indicate time when the detectors were on (white) off (gray) and times when the data was missing (MIA) blue.  The blue fraction should be << 1%
#
# ### PSDs
#
# \image html gstlal_llcbcnode03.png
#
# These represent the instantaneous PSD estimates of the running job.  The horizon distance is also computed.
#
# ### SNR / Chi-squared stats
#
# \image html gstlal_llcbcnode04.png
#
# These represent the instantaneous, cumulative SNR/chi-squared statistics for the job as well as the likelihood ranking plot.
#


def to_png_image(matplotlib_figure = None):
	# matplotlib_figure is needed because plot_likelihood_ratio_pdf
	# generates plots using matplotlib.figure while other plotting
	# functions use matplotlib.pyplot
	f = StringIO.StringIO()
	if matplotlib_figure is None:
		plt.savefig(f, format="png")
	else: 
		matplotlib_figure.savefig(f, format="png")
	print '<img src="data:image/png;base64,',base64.b64encode(f.getvalue()),'"></img>'
	f.close()

def plot(dataurl, plottype, xlabel = "", ylabel = "", title = "", textbox = None):
	fig = plt.figure(figsize=(5,3.5),)
	fig.patch.set_alpha(0.0)
	h = fig.add_subplot(111, axisbg = 'k')
	plt.subplots_adjust(bottom = 0.2, left = .16)
	plt.grid(color=(0.1,0.4,0.5), linewidth=2)
	try:
		data = numpy.loadtxt(dataurl)
	except IOError:
		print "Data could not be loaded"
		return
	except ValueError:
		print "<h2>Data is mis formatted, perhaps the webserver returned an error</h2>"
		return

	if plottype == "history":
		h.semilogy(data[:,0] - data[-1,0], data[:,1], 'w', alpha=0.75, linewidth=2)
		plt.ylim([min(data[:,1]), max(data[:,1])])
		locs = [min(data[:,1]), numpy.median(data[:,1]), max(data[:,1])]
		labels = ['%.2g' % lab for lab in locs]
		plt.yticks(locs, labels)
	elif plottype == "plot":
		h.plot(data[:,0], data[:,1], 'w', alpha=0.75, linewidth=2)
	elif plottype == "hist":
		from scipy.interpolate import interp1d
		x = numpy.linspace(data[0,0], data[-1,0], 100)
		f = interp1d(data[:,0], data[:,1], kind='linear')
		y = f(x)
		y[y<0] = 0.
		h.fill_between(x, y,  alpha=0.75, linewidth=2, facecolor="w", color="w")
	else:
		raise ValueError(plottype)

	if textbox is not None:
		plt.figtext(0.5, 0.2, textbox, color="w")

	plt.xlabel(xlabel)
	plt.ylabel(ylabel)
	plt.title(title)
	plt.savefig(sys.stdout, format="svg")


def ceil10(x):
	return 10**math.ceil(math.log10(x))

def floor10(x):
	return 10**math.floor(math.log10(x))


def psdplot(instrument, psd, fmin = 10., fmax = 2048.):
	fig = plt.figure(figsize=(5,3.5),)
	fig.patch.set_alpha(0.0)
	h = fig.add_subplot(111, axisbg = 'k')
	plt.subplots_adjust(bottom = 0.2, left = .16)
	plt.grid(color=(0.1,0.4,0.5), linewidth=2)

	y = psd.data
	f = numpy.linspace(psd.f0, len(y) * psd.deltaF, len(y)) + psd.f0

	print len(y), len(f)

	h.loglog(f, y, 'w', alpha=0.75, linewidth=2)

	plt.figtext(0.5, 0.2, "Horizon %.0f Mpc" % reference_psd.horizon_distance(psd, m1 = 1.4, m2 = 1.4, snr = 8., f_min = 30.), color="w")

	imin = bisect.bisect_left(f, fmin)
	imax = bisect.bisect_right(f, fmax)
	ymax = ceil10(y[imin:imax].max())
	ymin = max(floor10(y[imin:imax].min()), ymax / 1e6)
	plt.xlim((fmin, fmax))
	#plt.ylim((ymin, ymax))
	plt.ylim((1e-50, 1e-40))

	plt.xlabel("Frequency (Hz)")
	plt.ylabel("Power (strain^2/Hz)")
	plt.title(instrument)
	plt.savefig(sys.stdout, format="svg")


def livetime_plot(disconturl, livetimeurl, xlabel = "", ylabel = "", title = ""):
	fig = plt.figure(figsize=(5,3),)
	fig.patch.set_alpha(0.0)
	h = fig.add_subplot(111, axisbg = 'k', aspect='equal')
	plt.subplots_adjust(bottom = 0, left = .25, top = 1, right = .75)
	plt.grid(color="w")

	try:	
		livetimedata = numpy.loadtxt(livetimeurl)
		discontdata = numpy.loadtxt(disconturl)
	except IOError:
		print "Data could not be loaded"
		return

	dt = livetimedata[2]
	lt = livetimedata[1]
	discont = discontdata[1]
	# FIXME Hack to adjust for high sample rate L1 and H1 state vector
	if "V1" not in title:
		dt /= 16
		lt /= 16
		discont /= 16
	data = [dt, lt, discont]
	explode = [0.0, 0, 0.15]
	labels = ["OFF : %g (s)" % dt, "ON : %g (s)" % lt, "MIA : %g (s)" % discont]

	h.pie(data, shadow=True, explode = explode, labels = labels, autopct='%1.1f%%', colors = ('0.5', '1.0', (0.7, 0.7, 1.)))

	plt.xlabel(xlabel)
	plt.ylabel(ylabel)
	plt.title(title)
	plt.savefig(sys.stdout, format="svg")

# FIXME plot_likelihood_ratio_pdf is copied (with slight modifications) from
# gstlal_inspiral_plot_background. We considered moving plotting functions from
# gstlal_inspiral_plot_background to a python package plot_background.py and
# importing it instead of repeating the function definitions, but this is the
# only function we're using in here (for now) so making a seperate file doesn't
# make sense at the moment
def plot_likelihood_ratio_pdf(instruments, pdf, (xlo, xhi), tag, zerolag_pdf, size_inches):
	fig = figure.Figure()
	FigureCanvas(fig)
	fig.set_size_inches(size_inches)
	axes = fig.gca()
	axes.set_position((.15, .14, .84, .76))
	axes.semilogy(pdf.bins[0].centres(), pdf.array, color = "k", label="Background")
	if zerolag_pdf is not None:
		axes.semilogy(zerolag_pdf.bins[0].centres(), zerolag_pdf.array, color = "k", linestyle = "--", label="Zero-lag")
	axes.grid(which = "both")
	if instruments is None:
		axes.set_title(r"%s Log Likelihood Ratio PDF" % tag)
	else:
		axes.set_title(r"%s %s Log Likelihood Ratio PDF" % (", ".join(sorted(instruments)), tag))
	# FIXME latex text throws a warning about matplotlib missing fonts and
	# then doesn't display latex symbols correctly
	#axes.set_xlabel(r"$\ln \Lambda$")
	#axes.set_ylabel(r"$P(\ln \Lambda | \mathrm{%s})$" % tag.lower())
	# FIXME ylabel assumes all plots have noise tag
	axes.set_xlabel("ln Lambda")
	axes.set_ylabel("P(ln Lambda | n)")
	yhi = pdf[xlo:xhi,].max()
	ylo = pdf[xlo:xhi,].min()
	if zerolag_pdf is not None:
		yhi = max(yhi, zerolag_pdf[xlo:xhi,].max())
		ylo = min(ylo, zerolag_pdf[xlo:xhi,].min())
	ylo = max(yhi * 1e-40, ylo)
	axes.set_ylim((10**math.floor(math.log10(ylo) - .5), 10**math.ceil(math.log10(yhi) + .5)))
	axes.set_xlim((xlo, xhi))
	axes.legend()
	return fig


if "dir" not in form:
	raise ValueError("must specify dir")
if "id" not in form:
	raise ValueError("must specify id")

baseurl = '%s/%s' % (form.getvalue("dir"), form.getvalue("id"))

print >>sys.stdout, 'Cache-Control: no-cache, must-revalidate'
print >>sys.stdout, 'Expires: Mon, 26 Jul 1997 05:00:00 GMT'
print >>sys.stdout, 'Content-type: text/html\r\n'

print """
<html>
<head>
<meta http-equiv="Pragma" content="no-cache">
<meta http-equiv="Expires" content="-1">
<meta http-equiv="CACHE-CONTROL" content="NO-CACHE">
<meta http-equiv="refresh" content="300">
  <link rel="stylesheet" href="//code.jquery.com/ui/1.10.0/themes/base/jquery-ui.css" />
  <script src="//code.jquery.com/jquery-1.8.3.js"></script>
  <script src="//code.jquery.com/ui/1.10.0/jquery-ui.js"></script>
  <script type="text/javascript"> $(function() {
    $("#accordion").accordion({
    });

  });</script>
</head>
<body>
<img src="../lcbo.jpg"> <font size=10>Low-latency Compact Binary Online %d </font><hr><br>
""" % int(lal.GPSTimeNow())

print '<div id="accordion">'

print '<h1>Trigger stats</h1><div>'
plot(baseurl+'/latency_history.txt', "history", ylabel = "Latency (s)", xlabel = "Time since last trigger (s)")
plot(baseurl+'/latency_histogram.txt', "hist", ylabel = "Count", xlabel = "Latency (s)")
plot(baseurl+'/snr_history.txt', "history", ylabel = "SNR", xlabel = "Time since last trigger (s)")
plot(baseurl+'/ram_history.txt', "history", ylabel = "RAM (GB)", xlabel = "Time before now (s)")
print "</div>"

# Pie table
print "<h1> Live time </h1><div>"
for ifo in ("H1", "L1", "V1"):
	try:
		livetime_plot(baseurl+'/%s/strain_add_drop.txt' % ifo, baseurl+'/%s/state_vector_on_off_gap.txt' % ifo, title = ifo)
	except:
		print "could not get livetime data for ", ifo 
print"</div>"

# Psd table
print "<h1>PSDs</h1><div>"
try:
	psds = lalseries.read_psd_xmldoc(ligolw_utils.load_url("%s/psds.xml" % baseurl, contenthandler = LIGOLWContentHandler))
except IOError:
	print "Data could not be loaded"
except ValueError:
	print "<h2>Data is mis formatted, perhaps the webserver returned an error</h2>"
else:
	for instrument, psd in sorted(psds.items()):
		psdplot(instrument, psd)
print"</div>"

# likelihoods
path = form.getlist("url")[0]

coinc_params_distributions, ranking_data, seglists = far.parse_likelihood_control_doc(ligolw_utils.load_filename(path, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler))

print >> sys.stderr, "smooth" in form

if "smooth" in form:
	counts = coinc_params_distributions.background_pdf
	inj = coinc_params_distributions.injection_pdf
else:
	counts = coinc_params_distributions.background_rates
	inj = coinc_params_distributions.injection_rates

bgcol = (224/255.,224/255.,224/255.)

likely = copy.deepcopy(inj)
for i, ifo in enumerate(['H1','L1', 'V1']):
	print "<h1>%s SNR / CHISQ stats</h1>" % ifo
	print "<div>"
	likely[ifo+"_snr_chi"].array /= counts[ifo+"_snr_chi"].array
	for name, obj in (("background", counts), ("injections", inj), ("likelihood", likely)):
		fig = plt.figure(figsize=(6,5), facecolor = 'g')
		fig.patch.set_alpha(0.0)
		#plt.gray()
		H1 = obj[ifo+"_snr_chi"].array
		snr = obj[ifo+"_snr_chi"].bins[0].centres()[1:-1]
		chi = obj[ifo+"_snr_chi"].bins[1].centres()[1:-1]
		chi[0] = 0 # not inf
		ax = plt.subplot(111)
		plt.pcolormesh(snr, chi, numpy.log10(H1.T +1)[1:-1,1:-1])
		ax.set_yscale('log')
		plt.colorbar()
		plt.xlabel('SNR')
		plt.ylabel('reduced chi^2 / SNR^2')
		plt.ylim([.001, .5])
		plt.xlim([4,100])
		ax.set_xscale('log')
		plt.title('%s: %s log base 10 (number + 1)' % (ifo, name))
		plt.grid(color=(0.1,0.4,0.5), linewidth=2)
		to_png_image()
	print "</div>"

print "<h1> Likelihood Ratio PDFs </h1>"
print "<div>"

# Calculuate combined pdfs
ranking_data.finish()
for instruments, binnedarray in ranking_data.background_likelihood_pdfs.items() if ranking_data is not None else ():
	fig = plot_likelihood_ratio_pdf(instruments, binnedarray, (-5.,100.), "Noise", ranking_data.zero_lag_likelihood_pdfs[instruments], (6,5))
	to_png_image(matplotlib_figure = fig)
print "</div>"

#FIXME don't hard code CIT values
node = urlparse.urlparse(open("%s_registry.txt" % baseurl).readline()).netloc.split(":")[0]
print "<h1> Cluster node health<h1>"
print "<div>"
print "<img src=https://ldas-gridmon.ligo.caltech.edu/ganglia/graph.php?r=hour&z=xlarge&h=%s.cluster.ldas.cit&m=load_one&s=by+name&mc=2&g=cpu_report&c=Nodes/>" % node
print "<img src=https://ldas-gridmon.ligo.caltech.edu/ganglia/graph.php?r=hour&z=xlarge&h=%s.cluster.ldas.cit&m=load_one&s=by+name&mc=2&g=mem_report&c=Nodes/>" %node
print"</div>"
print "</body>"
