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

## @file gstlal_inspiral_lvalert_psd_plotter
# A program to listen to lvalerts, download the psd from gstlal gracedb events, plot it, and upload the results
#
# ### Command line interface
#
#	+ `--no-upload`: Write plots to disk.
#	+ `--skip-404`: Skip events that give 404 (file not found) errors (default is to abort).
#	+ `--verbose`: Be verbose.
#
# =============================================================================
#
#                                   Preamble
#
# =============================================================================
#


import httplib
import logging
import math
import numpy
from optparse import OptionParser
import os.path
import StringIO
import sys
import time
import urlparse
import json


from glue.ligolw import ligolw
from glue.ligolw import array as ligolw_array
from glue.ligolw import param as ligolw_param
from glue.ligolw import lsctables
from glue.ligolw import utils as ligolw_utils
from gstlal.reference_psd import horizon_distance
from gstlal import plotpsd
from ligo.gracedb import rest as gracedb
from pylal import series as lal_series


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


plotpsd.matplotlib.rcParams.update({
	"font.size": 8.0,
	"axes.titlesize": 10.0,
	"axes.labelsize": 10.0,
	"xtick.labelsize": 8.0,
	"ytick.labelsize": 8.0,
	"legend.fontsize": 8.0,
	"figure.dpi": 175,
	"savefig.dpi": 175,
	"text.usetex": True,
	"path.simplify": True
})


#
# =============================================================================
#
#                                   Library
#
# =============================================================================
#


def get_filename(gracedb_client, graceid, filename, retries = 3, retry_delay = 10.0, ignore_404 = False):
	for i in range(retries):
		logging.info("retrieving \"%s\" for %s" % (filename, graceid))
		response = gracedb_client.files(graceid, filename)
		if response.status == httplib.OK:
			return response
		if response.status == httplib.NOT_FOUND and ignore_404:
			logging.warning("retrieving \"%s\" for %s: (%d) %s.  skipping ..." % (filename, graceid, response.status, response.reason))
			return None
		logging.warning("retrieving \"%s\" for %s: (%d) %s.  pausing ..." % (filename, graceid, response.status, response.reason))
		time.sleep(retry_delay)
	raise Exception("retrieving \"%s\" for %s: (%d) %s" % (filename, graceid, response.status, response.reason))


def get_psds(gracedb_client, graceid, filename = "psd.xml.gz", ignore_404 = False):
	response = get_filename(gracedb_client, graceid, filename = filename, ignore_404 = ignore_404)
	if response is None:
		return None
	return lal_series.read_psd_xmldoc(ligolw_utils.load_fileobj(response, contenthandler = LIGOLWContentHandler)[0])


def get_coinc_xmldoc(gracedb_client, graceid, filename = "coinc.xml"):
	return ligolw_utils.load_fileobj(get_filename(gracedb_client, graceid, filename = filename), contenthandler = LIGOLWContentHandler)[0]


def upload_fig(fig, gracedb_client, graceid, filename = "psd.png"):
	plotfile = StringIO.StringIO()
	fig.savefig(plotfile, format = os.path.splitext(filename)[-1][1:])
	logging.info("uploading \"%s\" for %s" % (filename, graceid))
	response = gracedb_client.writeLog(graceid, "strain spectral density plot", filename = filename, filecontents = plotfile.getvalue(), tagname = "psd")
	if response.status != httplib.CREATED:
		raise Exception("upload of \"%s\" for %s failed: %s" % (filename, graceid, response["error"]))


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


def parse_command_line():
	parser = OptionParser()
	parser.add_option("--no-upload", action = "store_true", help = "Write plots to disk.")
	parser.add_option("--skip-404", action = "store_true", help = "Skip events that give 404 (file not found) errors (default is to abort).")
	parser.add_option("--gracedb-service-url", default="%s" % gracedb.DEFAULT_SERVICE_URL, help = "GraceDb service url to upload to (default: %s)" % gracedb.DEFAULT_SERVICE_URL)
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	options, graceids = parser.parse_args()

	if not graceids:
		# FIXME:  lvalert_listen doesn't allow command-line options
		options.verbose = True

	# can only call basicConfig once (otherwise need to switch to more
	# complex logging configuration)
	if options.verbose:
		logging.basicConfig(format = "%(asctime)s:%(message)s", level = logging.INFO)
	else:
		logging.basicConfig(format = "%(asctime)s:%(message)s")

	return options, graceids


#
# =============================================================================
#
#                                     Main
#
# =============================================================================
#


options, graceids = parse_command_line()


if not graceids:
	lvalert_data = json.loads(sys.stdin.read())
	logging.info("%(alert_type)s-type alert for event %(uid)s" % lvalert_data)
	logging.info("lvalert data: %s" % repr(lvalert_data))
	filename = os.path.split(urlparse.urlparse(lvalert_data["file"]).path)[-1]
	if filename not in (u"psd.xml.gz",):
		logging.info("filename is not 'psd.xml.gz'.  skipping")
		sys.exit()
	graceids = [str(lvalert_data["uid"])]


gracedb_client = gracedb.GraceDb(service_url = options.gracedb_service_url)


for graceid in graceids:
	psds = get_psds(gracedb_client, graceid, ignore_404 = options.skip_404)
	if psds is None:
		continue
	fig = plotpsd.plot_psds(psds, get_coinc_xmldoc(gracedb_client, graceid), plot_width = 1024)
	try:
		fig.tight_layout()
	except AttributeError:
		# not available on old matplotlibs
		pass
	if options.no_upload:
		filename = "psd_%s.png" % graceid
		logging.info("writing %s ..." % filename)
		fig.savefig(filename)
	else:
		upload_fig(fig, gracedb_client, graceid)
	logging.info("finished processing %s" % graceid)
