#!/usr/bin/env python
#
# Copyright (C) 2009-2011  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.
"""Stream-based inspiral analysis tool"""


#
# =============================================================================
#
#                                   Preamble
#
# =============================================================================
#


import os
import sys
from optparse import OptionParser
import resource
import signal
import subprocess
import time
import socket


# The following snippet is taken from http://gstreamer.freedesktop.org/wiki/FAQ#Mypygstprogramismysteriouslycoredumping.2Chowtofixthis.3F
import pygtk
pygtk.require("2.0")
import gobject
gobject.threads_init()
import pygst
pygst.require("0.10")
import gst


from glue import segments
from glue import segmentsUtils
from glue.ligolw import ligolw
from glue.ligolw import array
from glue.ligolw import param
from glue.ligolw import lsctables
array.use_in(ligolw.LIGOLWContentHandler)
param.use_in(ligolw.LIGOLWContentHandler)
lsctables.use_in(ligolw.LIGOLWContentHandler)
from glue.ligolw import utils
from glue.ligolw.utils import segments as ligolw_segments
from pylal.datatypes import LIGOTimeGPS
from pylal import series as lalseries
from gstlal import svd_bank
from gstlal import pipeparts
from gstlal import lloidparts
from gstlal import cbc_template_iir
from gstlal import far
from gstlal import inspiral
from gstlal import httpinterface
from gstlal import bottle
from pylal.date import XLALUTCToGPS


def excepthook(*args):
	# system exception hook that forces hard exit.  without this,
	# exceptions that occur inside python code invoked as a call-back
	# from the gstreamer pipeline just stop the pipeline, they don't
	# cause gstreamer to exit.

	# FIXME:  they probably *would* cause if we could figure out why
	# element errors and the like simply stop the pipeline instead of
	# crashing it, as well.  Perhaps this should be removed when/if the
	# "element error's don't crash program" problem is fixed
	sys.__excepthook__(*args)
	os._exit(1)

sys.excepthook = excepthook

# Setup the resource limits

maxproc = resource.getrlimit(resource.RLIMIT_NPROC)[1]
resource.setrlimit(resource.RLIMIT_NPROC, (maxproc, maxproc))
maxstack = resource.getrlimit(resource.RLIMIT_STACK)[1]
resource.setrlimit(resource.RLIMIT_STACK, (1 * 1024**2, maxstack)) # 1MB per thread, not 10


#
# disable time stamp checking
#


pipeparts.mkchecktimestamps = lambda pipeline, src, *args: src


# FIXME repair the shared memory partition just in case, update this as appropriate
def smrepair():
	subprocess.call(["smrepair", "LHO_Data"])
	subprocess.call(["smrepair", "LLO_Data"])
	subprocess.call(["smrepair", "VIRGO_Data"])


def now():
	return XLALUTCToGPS(time.gmtime())
	#return LIGOTimeGPS(time.time() - 315964786)


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


def parse_command_line():
	parser = OptionParser(
		description = __doc__
	)
	parser.add_option("--channel-name", metavar = "name", default = [], action = "append", help = "Set the name of the channel to process (optional).  The default is \"LSC-STRAIN\" for all detectors. Override with IFO=CHANNEL-NAME can be given multiple times")
	parser.add_option("--output", metavar = "filename", help = "Set the name of the LIGO light-weight XML output file *.{xml,xml.gz} or an SQLite database *.sqlite. If not given uses sqlite with format T050017-00")
	parser.add_option("--reference-psd", metavar = "filename", help = "Instead of measuring the noise spectrum, load the spectrum from this LIGO light-weight XML file (optional).")
	parser.add_option("--track-psd", action = "store_true", help = "Track PSD even if a reference is given")
	parser.add_option("--iir-bank", metavar = "filename", help = "Set the name of the LIGO light-weight XML file from which to load the iir bank for a given instrument in the form ifo:file These can be given as a comma separated list such as H1:file1,H2:file2,L1:file3 to analyze multiple instruments.")
        parser.add_option("--svd-bank", metavar = "filename", help = "Set the name of the LIGO light-weight XML file from which to load the svd bank for a given instrument in the form ifo:file These can be given as a comma separated list such as H1:file1,H2:file2,L1:file3 to analyze multiple instruments.")
	parser.add_option("--time-slide-file", metavar = "filename", help = "Set the name of the xml file to get time slide offsets")
	parser.add_option("--control-peak-time", metavar = "time", type = "int", help = "Set a time window in seconds to find peaks in the control signal")
	parser.add_option("--fir-stride", metavar = "time", type = "int", default = 8, help = "Set the length of the fir filter stride in seconds. default = 8")
	parser.add_option("--thinca-interval", metavar = "secs", type = "float", default = 30.0, help = "Set the thinca interval, default = 30s")
	parser.add_option("--ht-gate-threshold", metavar = "threshold", type = "float", help = "Set the threshold on whitened h(t) to mark samples as gaps (glitch removal)")
	parser.add_option("--coincidence-threshold", metavar = "value", type = "float", default = 0.020, help = "Set the coincidence window in seconds (default = 0.020).  The light-travel time between instruments will be added automatically in the coincidence test.")
	parser.add_option("--likelihood-file", metavar = "filename", help = "Set the name of the file from which to load initial likelihood ratio data (required).")
	parser.add_option("--marginalized-likelihood-file", metavar = "filename", help = "Set the name of the file from which to load initial marginalized likelihood ratio data (required).")
	parser.add_option("--write-pipeline", metavar = "filename", help = "Write a DOT graph description of the as-built pipeline to this file (optional).  The environment variable GST_DEBUG_DUMP_DOT_DIR must be set for this option to work.")
	parser.add_option("--comment", help = "Set the string to be recorded in comment and tag columns in various places in the output file (optional).")
	parser.add_option("--job-tag", help = "Set the string to identify this job and register the resources it provides on a node. Should 4 digits of the form 0001, 0002, etc.  required")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose (optional).")
	parser.add_option("-t", "--tmp-space", metavar = "path", help = "Path to a directory suitable for use as a work area while manipulating the database file.  The database file will be worked on in this directory, and then moved to the final location when complete.  This option is intended to improve performance when running in a networked environment, where there might be a local disk with higher bandwidth than is available to the filesystem on which the final output will reside.")
	parser.add_option("--gracedb-far-threshold", type = "float", help = "false alarm rate threshold for gracedb (Hz), if not given gracedb events are not sent")
	parser.add_option("--gracedb-type", default = "LowMass", help = "gracedb type, default is LowMass")
	parser.add_option("--gracedb-group", default = "Test", help = "gracedb group, default is Test")
	parser.add_option("--fake-data", metavar = "[white|silence|AdvVirgo|LIGO|AdvLIGO]", help = "Instead of reading data from .gwf files, generate and process coloured Gaussian noise.")
	parser.add_option("--injections", metavar = "filename", help = "Set the name of the LIGO light-weight XML file from which to load injections (optional).")
	parser.add_option("--veto-segments-file", metavar = "filename", help = "Set the name of the LIGO light-weight XML file from which to load vetoes (optional).")
	parser.add_option("--veto-segments-name", metavar = "name", help = "Set the name of the segments to extract from the segment tables and use as the veto list.", default = "vetoes")
	parser.add_option("--state-vector-on-bits", metavar = "name", default = [], action = "append", help = "Set the state vector on bits to process (optional).  The default is 0x7 for all detectors. Override with IFO=bits can be given multiple times")
	parser.add_option("--state-vector-off-bits", metavar = "name", default = [], action = "append", help = "Set the state vector off bits to process (optional).  The default is 0x160 for all detectors. Override with IFO=bits can be given multiple times")

	options, filenames = parser.parse_args()


	if options.reference_psd is None and not options.track_psd:
		raise ValueError("must use --track-psd if no reference psd is given, you can use both simultaneously")

	required_options = ["likelihood_file", "job_tag"]

	if options.svd_bank is not None:
		options.iir_bank = options.svd_bank
	
	missing_options = []
	if options.iir_bank is None:
		missing_options += ["--iir-bank"]
	missing_options += ["--%s" % option.replace("_", "-") for option in required_options if getattr(options, option) is None]
	if missing_options:
		raise ValueError, "missing required option(s) %s" % ", ".join(sorted(missing_options))

	# sanity check the job id
	if len(options.job_tag) != 4:
		raise ValueError("job id should be a 4 digit string that can convert to an int")	
	try:
		int(options.job_tag)
	except ValueError:
		raise ValueError("job id should be a 4 digit string that can convert to an int")

 	# Get the banks and make the detectors
	# FIXME add error checking on length of banks per detector, etc
	iir_banks = inspiral.parse_banks(options.iir_bank)
	
	# You get a dictionary of channels keyed by ifo, can be overidden by command line, default is LSC-STRAIN
	channel_dict = inspiral.channel_dict_from_channel_list(options.channel_name)
	detectors = {}
	for instrument in set(iir_banks.keys()):
		detectors[instrument] = lloidparts.DetectorData(None, channel_dict[instrument])

	# FIXME: should also check for read permissions
	required_files = []
	for instrument in iir_banks:
		required_files.extend(iir_banks[instrument])

	# autochisq
	options.chisq_type = "autochisq"

	# do this before converting option types
	process_params = options.__dict__.copy()

	options.state_vector_on_off_dict = inspiral.state_vector_on_off_dict_from_bit_lists(options.state_vector_on_bits, options.state_vector_off_bits)
	print options.state_vector_on_off_dict

	# hard-coded.  this number is needed in a few places, and storing
	# it in with the options is a convenient home
	options.psd_fft_length = 8	# seconds

	# this gets set so that if you log into a node you can find out what the job id is easily
	os.environ['GSTLAL_LL_JOB'] = options.job_tag

	#
	# Set up a registry of the resources that this job provides
	#
	
	host = socket.gethostname()

	# FIXME 
	# update this as bottle routes are added, can we do this automatically?
	fname = os.path.join(os.getcwd(), os.environ['GSTLAL_LL_JOB'] + "_registry.txt")
	f = open(fname, "w")

	@bottle.route("/registry.txt")
	def register(fname = None):
		
		yield "# %s %s\n" % (options.job_tag, host)

		# First do urls that do not depend on instruments
		for request in ("registry.txt", "gracedb_far_threshold.txt", "latency_histogram.txt", "latency_history.txt", "snr_history.txt", "ram_history.txt", "likelihood.xml", "bank.txt"):
			# FIXME don't hardcode port number
			yield "http://%s:16953/%s\n" % (host, request)

		# Then do instrument dependent urls
		for ifo in set(iir_banks.keys()):
			for request in ("strain_add_drop.txt", "state_vector_on_off_gap.txt", "psd.txt"):
				# FIXME don't hardcode port number
				yield "http://%s:16953/%s/%s\n" % (host, ifo, request)

	[f.write(l) for l in register(fname)]
	f.close()
	
	return options, filenames, process_params, iir_banks, detectors


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


#
# parse command line
#


options, filenames, process_params, iir_banks, detectors = parse_command_line()
frame_segments = dict([(instrument, None) for instrument in detectors])
if options.veto_segments_file is not None:
	veto_segments = ligolw_segments.segmenttable_get_by_name(utils.load_filename(options.veto_segments_file, verbose = options.verbose, contenthandler = ligolw.LIGOLWContentHandler), options.veto_segments_name).coalesce()
else:
	veto_segments = None


#
# set up the PSDs
#
# There are three modes for psds in this program
# 1) --reference-psd without --track-psd - a fixed psd (provided by the user) will be used to whiten the data
# 2) --track-psd without --reference-psd - a psd will me measured and used on the fly
# 3) --track-psd with --reference-psd - a psd will be measured on the fly, but the first "guess will come from the users provided psd
#


if options.reference_psd is not None:
	psd = lalseries.read_psd_xmldoc(utils.load_filename(options.reference_psd, verbose = options.verbose, contenthandler = ligolw.LIGOLWContentHandler))
else:
	psd = dict([(instrument, None) for instrument in detectors])


#
# Get IIR banks
#
# FIXME: snr_threshold
snr_threshold = 4.0
banks = {}

for instrument, files in iir_banks.items():
	for filename in files:
		bank = cbc_template_iir.load_iirbank(filename, snr_threshold, options.verbose)
		banks.setdefault(instrument,[]).append(bank)

for instrument in banks:
	for n, bank in enumerate(banks[instrument]):
		bank.logname = "%sbank%d" % (instrument,n)

#
# Build pipeline
#


if options.verbose:
	print >>sys.stderr, "assembling pipeline ...",

pipeline = gst.Pipeline("gstlal_inspiral")
mainloop = gobject.MainLoop()

# custom handler
class Handler(object):
	def __init__(self, mainloop, pipeline):
		self.mainloop = mainloop
		self.pipeline = pipeline

		bus = pipeline.get_bus()
		bus.add_signal_watch()
		bus.connect("message", self.on_message)

	def on_message(self, bus, message):
		if message.type == gst.MESSAGE_EOS:
			print >> sys.stderr, "EOS received"
			self.pipeline.set_state(gst.STATE_PAUSED)
			# FIXME we should set null, but if we try it locks
			# self.pipeline.set_state(gst.STATE_NULL)
			self.mainloop.quit()
		elif message.type == gst.MESSAGE_INFO:
			gerr, dbgmsg = message.parse_info()
			print >>sys.stderr, "info (%s:%d '%s'): %s" % (gerr.domain, gerr.code, gerr.message, dbgmsg)
		elif message.type == gst.MESSAGE_WARNING:
			gerr, dbgmsg = message.parse_warning()
			print >>sys.stderr, "warning (%s:%d '%s'): %s" % (gerr.domain, gerr.code, gerr.message, dbgmsg)
		elif message.type == gst.MESSAGE_ERROR:
			gerr, dbgmsg = message.parse_error()
			self.pipeline.set_state(gst.STATE_NULL)
			self.mainloop.quit()
			sys.exit("error (%s:%d '%s'): %s" % (gerr.domain, gerr.code, gerr.message, dbgmsg))

handler = Handler(mainloop, pipeline)

if options.fake_data is not None:
	seekevent = gst.event_new_seek(1., gst.FORMAT_TIME, gst.SEEK_FLAG_FLUSH | gst.SEEK_FLAG_KEY_UNIT, gst.SEEK_TYPE_SET, now().ns(), gst.SEEK_TYPE_SET, 2000000000000000000) # FIXME future infinity :)
else:
	seekevent = None

triggersrc = lloidparts.mkSPIIRmulti(
	pipeline,
	seekevent = seekevent, # seekevent iff fake data
	detectors = detectors,
	banks = banks,
	psd = psd,
	psd_fft_length = options.psd_fft_length,
	data_source = options.fake_data or "lvshm",
	injection_filename = options.injections,
	ht_gate_threshold = options.ht_gate_threshold,
	veto_segments = veto_segments,
	verbose = options.verbose,
	nxydump_segment = None,
	frame_segments = frame_segments,
	chisq_type = options.chisq_type,
	track_psd = options.track_psd,
	state_vector_on_off_dict = options.state_vector_on_off_dict
)


if options.verbose:
	print >>sys.stderr, "done"


#
# Load likelihood ratio data
# Note: It assumes injections are present
#


FAR, proc_id = far.LocalRankingData.from_xml(utils.load_filename(options.likelihood_file, verbose = options.verbose, contenthandler = ligolw.LIGOLWContentHandler))


#
# build output document
#


# FIXME now to "infinity"
start = now()
likelihood_snapshot_interval = 3600.0 # seconds
options.seg = segments.segment(LIGOTimeGPS(0), LIGOTimeGPS(2000000000))

if options.verbose:
	print >>sys.stderr, "initializing output document ..."
output = inspiral.Data(
	filename = options.output,
	process_params = process_params,
	pipeline = pipeline,
	instruments = set(detectors),
	seg = options.seg,
	injection_filename = None,
	time_slide_file = options.time_slide_file,
	coincidence_threshold = options.coincidence_threshold,
	FAR = FAR,
	marginalized_likelihood_file = options.marginalized_likelihood_file,	
	assign_likelihoods = True,
	likelihood_snapshot_interval = likelihood_snapshot_interval,	# seconds
	comment = options.comment,
	tmp_path = options.tmp_space,
	thinca_interval = options.thinca_interval,
	gracedb_far_threshold = options.gracedb_far_threshold,
	gracedb_type = options.gracedb_type,
	gracedb_group = options.gracedb_group,
	likelihood_file = options.likelihood_file,
	replace_file = True,# Replace output files that exist in online mode (none should exist if the files are autogenerated since they go by time stamp)
	verbose = options.verbose
)

if options.verbose:
	print >>sys.stderr, "... output document initialized"

if options.verbose:
	print >>sys.stderr, "attaching appsinks to pipeline ...",
appsync = pipeparts.AppSync(appsink_new_buffer = output.appsink_new_buffer)
appsinks = set(appsync.add_sink(pipeline, pipeparts.mkqueue(pipeline, src), caps = gst.Caps("application/x-lal-snglinspiral")) for src in triggersrc)
if options.verbose:
	print >>sys.stderr, "attached %d, done" % len(appsinks)


#
# if we request a dot graph of the pipeline, set it up
#


if options.write_pipeline is not None:
	pipeparts.connect_appsink_dump_dot(pipeline, appsinks, options.write_pipeline, options.verbose)
	pipeparts.write_dump_dot(pipeline, "%s.%s" % (options.write_pipeline, "NULL"), verbose = options.verbose)


#
# start http interface
#

# FIXME: don't hard-code port
httpinterface.start_servers(16953, verbose = options.verbose)


#
# Run pipeline
#


# repair the shared memory just in case before starting
smrepair()

if options.verbose:
	print >>sys.stderr, "setting pipeline state to playing ..."
if pipeline.set_state(gst.STATE_PLAYING) != gst.STATE_CHANGE_SUCCESS:
	raise RuntimeError, "pipeline did not enter playing state"

if options.write_pipeline is not None:
	pipeparts.write_dump_dot(pipeline, "%s.%s" % (options.write_pipeline, "PLAYING"), verbose = options.verbose)

if options.verbose:
	print >>sys.stderr, "running pipeline ..."


#
# setup sigint handler to shutdown pipeline
#

class SigData(object):
	def __init__(self):
		self.has_been_signaled = False

sigdata = SigData()

def signal_handler(signal, frame, pipeline = pipeline, instruments = detectors.keys(), sigdata = sigdata):
	if not sigdata.has_been_signaled:
		print >>sys.stderr, "*** SIG %d attempting graceful shutdown... ***" % (signal,)
		# override file name with approximate interval
		bus = pipeline.get_bus()
		bus.post(gst.message_new_eos(pipeline))
		sigdata.has_been_signaled = True
	else:
		print >>sys.stderr, "*** received SIG %d, but already handled... ***" % (signal,)

def crash_handler(signal, frame, pipeline = pipeline, instruments = detectors.keys(), sigdata = sigdata, mainloop = mainloop):
	if not sigdata.has_been_signaled:
		print >>sys.stderr, "*** SIG %d trying to copy outputfile, won't be able to shutdown pipeline***" % (signal,)
		mainloop.quit()
	else:
		print >>sys.stderr, "*** received SIG %d, but already handled... ***" % (signal,)

# this is how the program stops gracefully, it is the only way to stop it.
# Otherwise it runs forever man.

signal.signal(signal.SIGBUS, crash_handler)
signal.signal(signal.SIGSEGV, crash_handler)
signal.signal(signal.SIGABRT, crash_handler)
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)


#
# run pipeline
#


mainloop.run()


#
# done
#


# if no output is given use Naming Convention for LIGO Trigger Files T050017-00
output.write_output_file(filename = options.output or output.coincs_document.T050017_filename("%s_LLOID" % options.job_tag, "sqlite"), likelihood_file = options.likelihood_file, verbose = options.verbose)
# repair the shared memory just in case before exiting
smrepair()

# always end with an error so condor won't mark it as done
sys.exit(1)
