#!/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 sys
import os
os.environ["MPLCONFIGDIR"] = "/tmp"
import matplotlib
matplotlib.use('Agg')
import pylab
import numpy
import cgi
import cgitb
cgitb.enable()
form = cgi.FieldStorage()
import StringIO
import base64
import glob
import bisect
from glue import lal, segments
from glue.ligolw import utils
from glue.ligolw.utils import segments as ligolw_segments
from glue.ligolw import lsctables

from pylal.xlal.datatypes.ligotimegps import LIGOTimeGPS

def to_png_image():
        f = StringIO.StringIO()
        pylab.savefig(f, format="png")
        print '<img src="data:image/png;base64,',base64.b64encode(f.getvalue()),'"></img>'
        f.close()

def segment_table_to_state(segtable):
	time = {}
	state = {}
	for ifo in segtable:
		for seg in segtable[ifo]:
			time.setdefault(ifo, []).append(float(seg[0]))
			time.setdefault(ifo, []).append(float(seg[0]))
			time.setdefault(ifo, []).append(float(seg[1]))
			time.setdefault(ifo, []).append(float(seg[1]))
			state.setdefault(ifo, []).append(0)
			state.setdefault(ifo, []).append(1)
			state.setdefault(ifo, []).append(1)
			state.setdefault(ifo, []).append(0)
				
	return time, state

def parse_segment_files_from_cache(segcache):
	segdict = segments.segmentlistdict()
	for c in segcache:
		xmldoc = utils.load_url(c.url)
		segdef = lsctables.table.get_table(xmldoc, lsctables.SegmentDefTable.tableName)
		names = [r.name for r in segdef]
		names = set(names)
		for name in names:
			yield name, segment_table_to_state(ligolw_segments.segmenttable_get_by_name(xmldoc, name).coalesce())

def parse_segment_file(path, jobid, seg):
	segment_cache = lal.Cache.from_urls(glob.glob(os.path.join(path, '*_SEGMENTS-*.xml.gz')))
	return segment_cache.sieve(description = jobid, segment = seg)


def parse_id_and_duration(path, mchirp):
	data = []
	for fname in glob.glob(path + '/*_bank.txt'):
		try:
			idstr = os.path.split(fname)[1][0:4]
			f = open(fname)
			duration, mc = f.readline().split()[1:]
			data.append((float(mc), float(duration), idstr))
			f.close()
		except OSError:
			pass
	data.sort()

	# Identify the id for the requested chirp mass
	index = bisect.bisect_left([d[0] for d in data], mchirp) - 1 #because of how bisect works
	idstr = data[index][2]
	duration = float(data[index][1])
	return idstr, duration


path = form.getvalue("dir")
eventtime = float(form.getvalue("gps"))
mchirp = float(form.getvalue("mchirp"))
ifos = form.getlist("ifo")

idstr, duration = parse_id_and_duration(path, mchirp)
eventsegment = segments.segment((LIGOTimeGPS(eventtime-duration), LIGOTimeGPS(eventtime+duration)))

segment_cache = parse_segment_file(path, idstr, eventsegment)

times, states = parse_segment_files_from_cache(segment_cache)

#logfiles = glob.glob(os.path.join(path, 'logs/gstlal_inspiral-%s*.err' % (idstr,)))

#time, state = parse_log(logfiles)

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'

# HTML preamble
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">
</head>
<body bgcolor=#E0E0E0 face="Arial">
<br>
"""

for (name, (times, states)) in parse_segment_files_from_cache(segment_cache):
	for ifo in times:
		if ifos and ifo not in ifos:
			continue
		pylab.figure(figsize=(12,2))
		pylab.subplots_adjust(bottom = 0.1, top = 0.85)
		pylab.fill_between(numpy.array(times[ifo]) - eventtime, states[ifo], alpha=0.5)#, ls = "steps-post")
		pylab.xlim([-duration, duration])
		pylab.title("%s %s state before event at %.0f with mchirp %.3f and duration %.0f (s)" % (ifo, name, eventtime, mchirp, duration))
		pylab.yticks([])
		to_png_image()

print """
</body>
</html>
"""
		
