#!/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 resource
import sys
from optparse import OptionParser
import signal
import subprocess
import socket
import time

# 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 gstlal import bottle

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 pylal.date import XLALUTCToGPS
from gstlal import svd_bank
from gstlal import pipeparts
from gstlal import lloidparts
from gstlal import simplehandler
from gstlal import far
from gstlal import inspiral
from gstlal import httpinterface
from gstlal import datasource

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

#
# Make sure we have sufficient resources
# We allocate far more memory than we need, so this is okay
#

# set the number of processes up to hard limit
maxproc = resource.getrlimit(resource.RLIMIT_NPROC)[1]
resource.setrlimit(resource.RLIMIT_NPROC, (maxproc, maxproc))

# set the total set size up to hard limit
maxas = resource.getrlimit(resource.RLIMIT_AS)[1]
resource.setrlimit(resource.RLIMIT_AS, (maxas, maxas))

# set the stack size per thread to be smaller
maxstack = resource.getrlimit(resource.RLIMIT_STACK)[1]
resource.setrlimit(resource.RLIMIT_STACK, (1 * 1024**2, maxstack)) # 1MB per thread, not 10


def now():
	return XLALUTCToGPS(time.gmtime())

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


def parse_command_line():
	parser = OptionParser(
		description = __doc__
	)

	# append all the datasource specific options
	datasource.append_options(parser)

	parser.add_option("--psd-fft-length", metavar = "s", default = 16, type = "int", help = "FFT length, default 16s")
	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("--nxydump-segment", metavar = "start:stop", default = ":", help = "Set the time interval to dump from nxydump elments (optional).  The default is \":\", i.e. dump all time.")
	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 (required).")
	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("--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("--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("--chisq-type", metavar = "type", default = "autochisq", help = "Choose the type of chisq computation to perform. Must be one of (autochisq|timeslicechisq). The default is autochisq.")
	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("--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("--check-time-stamps", action = "store_true", help = "Turn on time stamp checking")
	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("--blind-injections", metavar = "filename", help = "Set the name of an injection file that will be added to the data without saving the sim_inspiral_table or otherwise processing the data differently.  Has the effect of having hidden signals in the input data.  --injections must not be specified in this case")

	# Online options

	parser.add_option("--job-tag", help = "Set the string to identify this job and register the resources it provides on a node. Should be 4 digits of the form 0001, 0002, etc.  required")
	parser.add_option("--likelihood-file", metavar = "filename", help = "Set the name of the likelihood ratio data file to use (optional).  If not specified, likelihood ratios will not be assigned to coincs.")
	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("--gracedb-far-threshold", type = "float", help = "false alarm rate threshold for gracedb (Hz), if not given gracedb events are not sent")
	parser.add_option("--likelihood-snapshot-interval", type = "float", metavar = "seconds", help = "How often to reread the marginalized likelihoood data and snapshot the trigger files.")
	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("--thinca-interval", metavar = "secs", type = "float", default = 30.0, help = "Set the thinca interval, default = 30s")

	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")

	if options.blind_injections and options.injections:
		raise ValueError("must use only one of --blind-injections or --injections")

	required_options = []

	missing_options = []
	if options.svd_bank is None:
		missing_options += ["--svd-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)))

	# parse the datasource specific information and do option checking
	detectors = datasource.GWDataSourceInfo(options)
	if len(detectors.channel_dict) < 2:
		raise ValueError("only coincident searches are supported:  must process data from at least two antennae")

	# Get the banks and make the detectors
	# FIXME add error checking on length of banks per detector, etc
	svd_banks = inspiral.parse_banks(options.svd_bank)

	# FIXME: should also check for read permissions
	required_files = []
	for instrument in svd_banks:
		required_files.extend(svd_banks[instrument])
	if options.veto_segments_file:
		required_files += [options.veto_segments_file]
	missing_files = [filename for filename in required_files if not os.path.exists(filename)]
	if missing_files:
		raise ValueError, "files %s do not exist" % ", ".join("'%s'" % filename for filename in sorted(missing_files))

	if options.chisq_type not in ["autochisq", "timeslicechisq"]:
		raise ValueError, "--chisq-type must be one of (autochisq|timeslicechisq), given %s" % (options.chisq_type)

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

	options.nxydump_segment, = segmentsUtils.from_range_strings([options.nxydump_segment], boundtype = LIGOTimeGPS)

	# Online specific initialization
	# FIXME someday support other online sources
	if options.data_source == "lvshm":
		# check required options in this case
		required_options = ["likelihood_file", "job_tag", "marginalized_likelihood_file"]

		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))

		# make an "infinite" extent segment
		detectors.seg = segments.segment(LIGOTimeGPS(0), LIGOTimeGPS(2000000000))

		# 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 %s\n" % (options.job_tag, host, " ".join(set(svd_banks.keys())))

			# 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", "segments.xml"):
				# FIXME don't hardcode port number
				yield "http://%s:16953/%s\n" % (host, request)

			# Then do instrument dependent urls
			for ifo in set(svd_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()
	else:
		bad_options = []
		for option in ["likelihood_file", "job_tag", "marginalized_likelihood_file", "likelihood_snapshot_interval"]:
			if getattr(options, option) is not None:
				bad_options.append(option)
		if bad_options:
			raise ValueError("%s options should only be given for online running" % ",".join(bad_options))
	# we're done
	return options, filenames, process_params, svd_banks, detectors


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


#
# parse command line
#


options, filenames, process_params, svd_banks, detectors = parse_command_line()

if not options.check_time_stamps:
	pipeparts.mkchecktimestamps = lambda pipeline, src, *args: src


#
# Parse the vetos segments file(s) if provided
#


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.channel_dict)


#
# Parse template banks
#


banks = inspiral.parse_bank_files(svd_banks, verbose = options.verbose)
@bottle.route("/bank.txt")
def get_filter_length_and_chirpmass(banks = banks):
	bank = banks.values()[0][0] #FIXME maybe shouldn't just take the first ones
	yield '%.14g %.4g %.4g' % (float(now()), bank.filter_length, bank.sngl_inspiral_table[0].mchirp)


#
# Build pipeline
#


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

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


triggersrc = lloidparts.mkLLOIDmulti(
	pipeline,
	detectors = detectors,
	banks = banks,
	psd = psd,
	psd_fft_length = options.psd_fft_length,
	ht_gate_threshold = options.ht_gate_threshold,
	veto_segments = veto_segments,
	verbose = options.verbose,
	nxydump_segment = options.nxydump_segment,
	chisq_type = options.chisq_type,
	track_psd = options.track_psd,
	control_peak_time = options.control_peak_time,
	fir_stride = options.fir_stride,
	blind_injections = options.blind_injections
)


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


#
# Load likelihood ratio data, assume injections are present!
#


if options.likelihood_file is not None:
	FAR, proc_id = far.LocalRankingData.from_xml(utils.load_filename(options.likelihood_file, verbose = options.verbose, contenthandler = ligolw.LIGOLWContentHandler))
else:
	FAR = far.LocalRankingData(segments.segment(None, None), far.DistributionsStats())
	# initialize an empty trials table with the combinations of detectors
	# that this job will analyze
	FAR.trials_table.initialize_from_sngl_ifos(detectors.channel_dict.keys())


#
# build output document
#


if options.verbose:
	print >>sys.stderr, "initializing output document ..."
output = inspiral.Data(
	filename = options.output or "%s-%s_LLOID-%d-%d.xml.gz" % (lsctables.ifos_from_instrument_set(detectors.channel_dict.keys()).replace(",", ""), options.job_tag, int(detectors.seg[0]), int(abs(detectors.seg))),
	process_params = process_params,
	pipeline = pipeline,
	instruments = set(detectors.channel_dict),
	seg = detectors.seg or segments.segment(LIGOTimeGPS(0), LIGOTimeGPS(2000000000)), # online data doesn't have a segment so make it all possible time
	injection_filename = options.injections,
	time_slide_file = options.time_slide_file,
	coincidence_threshold = options.coincidence_threshold,
	FAR = FAR,
	likelihood_file = options.likelihood_file,
	marginalized_likelihood_file = options.marginalized_likelihood_file,
	likelihood_snapshot_interval = options.likelihood_snapshot_interval,	# seconds
	assign_likelihoods = options.likelihood_file is not None,
	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,
	verbose = options.verbose
)
if options.verbose:
	print >>sys.stderr, "... output document initialized"

# setup the pipeline hander with dq gates if in online mode
if options.data_source == "lvshm":
	gates = dict([("%s_state_vector_gate" % ifo, "state vector gate") for ifo in detectors.channel_dict])
else:
	gates = {}
handler = lloidparts.Handler(mainloop, pipeline, gates = gates, tag = options.job_tag, dataclass = output, verbose = options.verbose)

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)


#
# Run pipeline
#


if options.data_source != "lvshm":
	if options.verbose:
		print >>sys.stderr, "setting pipeline state to paused ..."
	if pipeline.set_state(gst.STATE_PAUSED) != gst.STATE_CHANGE_SUCCESS:
		raise RuntimeError, "pipeline did not enter paused state"
else:
	#
	# start http interface
	#

	# FIXME: don't hard-code port, don't use in this program right now since more than one job runs per machine
	httpinterface.start_servers(16953, verbose = options.verbose)
	#
	# setup sigint handler to shutdown pipeline.  this is how the program stops
	# gracefully, it is the only way to stop it.  Otherwise it runs forever
	# man.
	#


	class OneTimeSignalHandler(object):
		def __init__(self, pipeline):
			self.pipeline = pipeline
			self.count = 0

		def __call__(self, signum, frame):
			self.count += 1
			if self.count == 1:
				print >>sys.stderr, "*** SIG %d attempting graceful shutdown (this might take several minutes) ... ***" % signum
				try:
					#FIXME how do I choose a timestamp?
					self.pipeline.get_bus().post(inspiral.message_new_checkpoint(self.pipeline, timestamp=now().ns()))
					if not self.pipeline.send_event(gst.event_new_eos()):
						raise Exception("pipeline.send_event(EOS) returned failure")
				except Exception, e:
					print >>sys.stderr, "graceful shutdown failed: %s\naborting." % str(e)
					os._exit(1)
			else:
				print >>sys.stderr, "*** received SIG %d %d times... ***" % (signum, self.count)

	signal.signal(signal.SIGINT, OneTimeSignalHandler(pipeline))
	signal.signal(signal.SIGTERM, OneTimeSignalHandler(pipeline))
	# FIXME get rid of this someday: Repair the shared memory just in case before starting
	for partition in ("LHO_Data", "LLO_Data", "VIRGO_Data"):
		subprocess.call(["smrepair", partition])

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 ..."
mainloop.run()


#
# write output file
#


output.write_output_file(filename = options.output or output.coincs_document.T050017_filename("%s_LLOID" % options.job_tag, "xml.gz"), likelihood_file = options.likelihood_file, verbose = options.verbose)


#
# done
#

if options.data_source == "lvshm":
	sys.exit(1) # online pipeline always ends with an error code
