#!/usr/bin/env python
#
# Copyright (C) 2012  Kipp Cannon, Chad Hanna, Drew Keppel
#
# 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
# A program to request some followup data from a running gstlal_inspiral job based on gracedb submissions notified by lvalert

import sys, os
import logging
os.environ["MPLCONFIGDIR"] = "/tmp"
import matplotlib
matplotlib.use('Agg')
import pylab
import numpy
import copy
from gstlal import far
from gstlal import plotfar
from gstlal import lvalert_helper
import time
from glue.ligolw import ligolw, lsctables
from glue.ligolw import utils as ligolw_utils
from glue.ligolw.utils import process as ligolw_process
import urllib
import urlparse
import httplib
import subprocess
from ligo.gracedb import rest as gracedb
import glob
import json
from optparse import OptionParser

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

def parse_command_line():
	parser = OptionParser()
	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)

	options, gid_list = parser.parse_args()
	
	if len(gid_list) > 1:
		raise ValueError("%d graceids specified, no more than one allowed" % len(gid_list))

	if len(gid_list) == 0:
		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()
		gid = str(lvalert_data["uid"])
	else:
		gid = gid_list[0]

	return options, gid
	

def plot_snrchisq(snr_chisq_dict, coinc_param_distributions, plot_type):
	for i, ifo in enumerate(['H1', 'L1', 'V1']):
		snrm, chisqm = snr_chisq_dict[ifo]
		if snrm is None:
			continue
		fig = plotfar.plot_snr_chi_pdf(coinc_param_distributions, ifo, plot_type, 400, event_snr = snrm, event_chisq = chisqm)
		try:
			fig.tight_layout()
		except AttributeError:
			# not in old matplotlibs
			pass
		fname = 'gracedb/%s/%s_%s_%s_snrchi.png' % (gid, gid, ifo, plot_type)
		fig.savefig(fname)
		gracedb.writeLog(gid, "%s SNR/Chisq" % ifo, filename = fname, filecontents = open(fname).read(), tagname = {"background_pdf":"background", "injection_pdf":"background", "zero_lag_pdf":"background", "LR":"background"}[plot_type])


options, gid = parse_command_line()

try:
	os.mkdir("gracedb/" + gid)
except OSError:
	pass

gracedb = gracedb.GraceDb(options.gracedb_service_url)
xmldoc, checksum = lvalert_helper.get_coinc_xmldoc(gracedb, gid)
coinc_inspiral_row = lsctables.CoincInspiralTable.get_table(xmldoc)[0]
sngl_inspiral_table = lsctables.SnglInspiralTable.get_table(xmldoc)
coinc_table = lsctables.CoincTable.get_table(xmldoc)[0]
snr_chisq_dict = dict(((ifo,(None,None)) for ifo in ['H1','L1', 'V1']))
for r in sngl_inspiral_table:
	snr_chisq_dict[r.ifo] = (r.snr, r.chisq)
try:
	path, = ligolw_process.get_process_params(xmldoc, "gstlal_inspiral", "--likelihood-file")
except ValueError:
	path, = ligolw_process.get_process_params(xmldoc, "gstlal_inspiral", "--reference-likelihood-file")
marg_path, = ligolw_process.get_process_params(xmldoc, "gstlal_inspiral", "--marginalized-likelihood-file")
tag, = ligolw_process.get_process_params(xmldoc, "gstlal_inspiral", "--job-tag")
node, = [r.node for r in lsctables.ProcessTable.get_table(xmldoc) if r.program == "gstlal_inspiral"]

url = open("%s_registry.txt" % tag).readline().strip()


#
# Background plots
#

coinc_param_distributions = far.parse_likelihood_control_doc(ligolw_utils.load_filename(path, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler))[0]
ranking_data, seglistsdict = far.parse_likelihood_control_doc(ligolw_utils.load_filename(marg_path, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler))[1:]
ranking_data.finish()
# FIXME Is this livetime correct?
fapfar = far.FAPFAR(ranking_data, livetime = far.get_live_time(seglistsdict))

counts = coinc_param_distributions.background_rates
inj = coinc_param_distributions.injection_rates

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

likely = copy.deepcopy(inj)
for plot_type in ["background_pdf", "injection_pdf", "zero_lag_pdf", "LR"]:
	plot_snrchisq(snr_chisq_dict, coinc_param_distributions, plot_type)

fig = plotfar.plot_likelihood_ratio_ccdf(fapfar, (-5.,100.), "Noise", event_likelihood = coinc_table.likelihood)
fname = 'gracedb/%s/%s_likehoodratio_ccdf.png' % (gid, gid)
try:
	fig.tight_layout()
except AttributeError:
	# not in old matplotlibs
	pass
fig.savefig(fname)
gracedb.writeLog(gid, "Likelihood Ratio CCDF", filename = fname, filecontents = open(fname).read(), tagname = "background")

fig = plotfar.plot_horizon_distance_vs_time(coinc_param_distributions, (coinc_inspiral_row.get_end() - 14400.,coinc_inspiral_row.get_end()), 241)
fname = 'gracedb/%s/%s_horizon_distances.png' % (gid, gid)
fig.savefig(fname)
gracedb.writeLog(gid, "Horizon Distances", filename = fname, filecontents = open(fname).read(), tagname = "psd")
