#!/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
from pylal import rate
os.environ["MPLCONFIGDIR"] = "/tmp"
import matplotlib
matplotlib.use('Agg')
import pylab
import numpy
from glue.ligolw import utils
import copy
from gstlal import far
from gstlal import plotfar
import time
from glue.ligolw import ligolw, table, lsctables, param, array
array.use_in(ligolw.LIGOLWContentHandler)
param.use_in(ligolw.LIGOLWContentHandler)
lsctables.use_in(ligolw.LIGOLWContentHandler)
import urllib
import urlparse
import httplib
import subprocess
from ligo.gracedb import DEFAULT_SERVICE_URL
from ligo.gracedb.rest import GraceDb
import glob
import json
from optparse import OptionParser

def parse_command_line():
	parser = OptionParser()
	parser.add_option("--gracedb-service-url", default="%s" % DEFAULT_SERVICE_URL, help = "GraceDb service url to upload to (default: %s)" % 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 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 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)
		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(options.gracedb_service_url)
filename = get_filename(gracedb, gid, "coinc.xml", retries = 3, retry_delay = 10.0, ignore_404 = False)
xmldoc, checksum = utils.load_fileobj(filename, contenthandler = ligolw.LIGOLWContentHandler)
coinc_inspiral_row = lsctables.table.get_table(xmldoc, lsctables.CoincInspiralTable.tableName)[0]
sngl_inspiral_table = lsctables.table.get_table(xmldoc, lsctables.SnglInspiralTable.tableName)
coinc_table = lsctables.table.get_table(xmldoc, lsctables.CoincTable.tableName)[0]
process = lsctables.table.get_table(xmldoc, lsctables.ProcessTable.tableName)
params = lsctables.table.get_table(xmldoc, lsctables.ProcessParamsTable.tableName)
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)
for r in params:
	if r.program == "gstlal_inspiral" and r.param == "--likelihood-file":
		path = r.value
	elif r.program == "gstlal_inspiral" and r.param == "--reference-likelihood-file":
		path = r.value
	if r.program =="gstlal_inspiral" and r.param == "--marginalized-likelihood-file":
		marg_path = r.value
	if r.program == "gstlal_inspiral" and r.param == "--job-tag":
		tag = r.value
for r in process:
	if r.program == "gstlal_inspiral":
		node = r.node

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

#
# SNR / Time plot
#

#FIXME don't hardcode port number
t,snr = numpy.loadtxt(urllib.urlopen("%ssnr_history.txt" % url), unpack = True)
t -= coinc_inspiral_row.get_end()
pylab.figure()
pylab.semilogy(t, snr)
pylab.ylabel('SNR')
pylab.xlabel('Time from event (s)')
fname = 'gracedb/%s/%s_snrtime.png' % (gid, gid)
pylab.ylim([min(snr), max(snr)])
pylab.xlim([min(t), max(t)])
pylab.grid()
pylab.savefig(fname)
gracedb.writeLog(gid, "SNR vs time", filename = fname, filecontents = open(fname).read(), tagname = "background")

#
# SNR / Chisq plots
#

coinc_param_distributions = far.parse_likelihood_control_doc(utils.load_filename(path, contenthandler = far.ThincaCoincParamsDistributions.LIGOLWContentHandler))[0]
ranking_data, seglistsdict = far.parse_likelihood_control_doc(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)
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")

sys.exit()

#
# QScans
# FIXME currently broken
#

for r in sngl_inspiral_table:
	cachefile = "%s/%sframes.cache" % (gid, r.ifo)
	subprocess.call(["createframecache.pl", cachefile, "/scratch/llcache/llhoft/%s" % r.ifo])
	conffile = "%s/%sconfig.txt" % (gid, r.ifo)
	f = open(conffile , "w")
	channel = "%s:%s" % (r.ifo, r.channel)
	f.write("""[%s,%s]

{
   channelName:                 '%s'
   frameType:                   '%s_llhoft'
   sampleFrequency:             16384
   searchFrequencyRange:        [10 2048]
   #searchChirpRange:            [+0 +0]
   searchTimeRange:              64
   searchQRange:                [4 64]
   searchMaximumEnergyLoss:     0.2
   searchWindowDuration:        0.25
   whiteNoiseFalseRate:         1e-3
   plotTimeRanges:              10
   plotFrequencyRange:          []
   plotNormalizedEnergyRange:   [0 25.5]
   alwaysPlotFlag:               1
} """ % (channel, channel, channel, r.ifo))
	f.close()

	tstr = str(r.get_end())

	# actually run the scan
	subprocess.call(["dmt_wscan", str(r.get_end()), conffile, cachefile, gid, "1", "1"])

# upload some plots to gracedb
for f in glob.glob("%s/*_spectrogram_autoscaled.png" % gid):
	gracedb.writeLog(gid, "DMT omega scan", filename = f, filecontents = open(f).read(), tagname = "background")
	
sys.exit()
