#!/usr/bin/env python
#
# Copyright (C) 2011  Chad Hanna
#
# 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
# This program will make create a HTCondor DAG to automate the running of low-latency, online gstlal_inspiral jobs; see gstlal_ll_trigger_pipe

"""
This program makes a dag for a gstlal inspiral low latency pipeline
"""

__author__ = 'Chad Hanna <channa@caltech.edu>'

#
# import standard modules and append the lalapps prefix to the python path
#

import sys, os, copy, stat
import shutil

#
# import the modules we need to build the pipeline
#

from glue import pipeline
from glue import lal
import glue.ligolw.utils as utils
from optparse import OptionParser
from gstlal import inspiral
from gstlal import inspiral_pipe
from gstlal import dagparts
from gstlal import datasource

##
# ### Graph of the HTCondor DAG
#
# - gray boxes are optional and depend on the command line given
#
# @dot
# digraph G {
#       // graph properties
#
#       rankdir=LR;
#       compound=true;
#       node [shape=record fontsize=10 fontname="Verdana"];     
#       edge [fontsize=8 fontname="Verdana"];
#	gstlal_inspiral [URL="\ref gstlal_inspiral"];
#	gstlal_llcbcsummary [URL="\ref gstlal_llcbcsummary"];
#	gstlal_llcbcnode [URL="\ref gstlal_llcbcnode"];
#	gstlal_inspiral_marginalize_likelihoods_online [URL="\ref gstlal_inspiral_marginalize_likelihoods_online"];
#	lvalert_listen [style=filled, color=lightgrey, URL="https://www.lsc-group.phys.uwm.edu/daswg/docs/howto/lvalert-howto.html"];
# }
# @enddot
#
# ### Usage cases
#
# - Typical usage case
#
# ### Command line options
#
#	+ `--psd-fft-length` [int] (s): FFT length, default 16s
#	+ `--reference-psd` [filename]: Set the reference psd file
#	+ `--bank-cache` [filenames]: Set the bank cache files in format H1=H1.cache,H2=H2.cache, etc
#	+ `--channel` [IFO=CHANNEL-NAME]: 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
#	+ `--ht-gate-threshold` [float]: Set the h(t) gate threshold to reject glitches
#	+ `--do-iir-pipeline`: Run the iir pipeline instead of lloid")
#	+ `--max-jobs` [num]: Stop parsing the cache after reaching a certain number of jobs to limit what is submitted to the HTCondor pool
#	+ `--likelihood-file` [filename]: Set the likelihood file, required if --copy-likelihoods is used
#	+ `--marginalized-likelihood-file` [filename]: Set the marginalized likelihood file, required
#	+ `--control-peak-time` [int] (s): Set the control peak time, default 4
#	+ `--fir-stride` [int] (s): Set the fir bank stride, default 4
#	+ `--thinca-interval` [int] (s): Set the thinca interval, default 10
#	+ `--gracedb-far-threshold` [float] (Hz): False alarm rate threshold for gracedb (Hz). Set to -1 to disable uploads
#	+ `--gracedb-type` [str]: Set the gracedb type, default LowMass
#	+ `--gracedb-group` [str]: Set the gracedb group, default Test
#	+ `--data-source` [lvshm|framexmit]: Where to get the data from. Default lvshm
#	+ `--copy-likelihoods`: Copy the likelihood files from a seed, must give --likelihood-file: WARNING overwrites existing files
#	+ `--veto-segments-file` [filename]: Set the name of the LIGO light-weight XML file from which to load vetoes (optional)
#	+ `--veto-segments-name` [name]: Set the name of the segments to extract from the segment tables and use as the veto list. Default is "vetoes"
#	+ `--state-vector-on-bits` [IFO=bits]: 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
#	+ `--state-vector-off-bits` [IFO=bits]: 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
#	+ `--lvalert-listener-program` [program]: Set the programs to respond to lvalerts from this analysis, can be given multiple times
#


class lvalert_listen_job(inspiral_pipe.generic_job):
	"""
	A lvalert_listen_job
	"""
	def __init__(self, program, gracedb_group = "CBC", gracedb_type = "LowMass", progs = ("gstlal_inspiral_lvalert_psd_plotter", "gstlal_inspiral_followups_from_gracedb")):
		"""
		"""
		inspiral_pipe.generic_job.__init__(self, program)

		# produce the lvalert processor

		f = open("lvalert.sh", "w")
		f.write("#!/bin/bash \n")
		f.write('cat <&0 | tee ')
		for prog in progs:
			f.write(">(%s) " % dagparts.which(prog))
		f.close()
		os.chmod('lvalert.sh', os.stat('lvalert.sh').st_mode | stat.S_IEXEC)

		f = open("lvalert.ini", "w")
		#FIXME gracedb server code sets up nodes based on this convention
		f.write("[%s_%s]\n" % (gracedb_group.lower(), gracedb_type.lower()))
		f.write("executable=./lvalert.sh")
		f.close()


class lvalert_listen_node(pipeline.CondorDAGNode):
	"""
	lvalert_listen node
	"""
	def __init__(self, job, dag):
		pipeline.CondorDAGNode.__init__(self,job)
		self.add_var_opt("username", raw_input("lvalert username: "))
		self.add_var_opt("password", raw_input("lvalert password: "))
		self.add_var_opt("config-file", "lvalert.ini")
		dag.add_node(self)


#
# Parse the command line
#


def parse_command_line():
	parser = OptionParser(description = __doc__)
	parser.add_option("--psd-fft-length", metavar = "s", default = 16, type = "int", help = "FFT length, default 16s")
	parser.add_option("--reference-psd", metavar = "filename", help = "Set the reference psd file.")
	parser.add_option("--bank-cache", metavar = "filenames", help = "Set the bank cache files in format H1=H1.cache,H2=H2.cache, etc..")
	parser.add_option("--channel", 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("--ht-gate-threshold", metavar = "float", help = "Set the h(t) gate threshold to reject glitches", type="float")
	parser.add_option("--do-iir-pipeline", action = "store_true", help = "run the iir pipeline instead of lloid")
	parser.add_option("--max-jobs", metavar = "num", type = "int", help = "stop parsing the cache after reaching a certain number of jobs to limit what is submitted to the HTCondor pool")
	parser.add_option("--likelihood-file", help = "set the likelihood file, required if --copy-likelihoods is used")	
	parser.add_option("--marginalized-likelihood-file", help = "set the marginalized likelihood file, required")	
	parser.add_option("--control-peak-time", default = 4, metavar = "secs", help = "set the control peak time, default 4")
	parser.add_option("--fir-stride", default = 4, metavar = "secs", help = "set the fir bank stride, default 4")
	parser.add_option("--thinca-interval", default = 10, metavar = "secs", help = "set the thinca interval, default 10")
	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 LowMass")
	parser.add_option("--gracedb-group", default = "Test", help = "gracedb group, default Test")
	parser.add_option("--data-source", metavar = "[lvshm|]", default = "lvshm", help = "Where to get the data from. Default lvshm")
	parser.add_option("--copy-likelihoods", action = "store_true", help = "Copy the likelihood files from a seed, must give --likelihood-file: : WARNING overwrites existing files")
	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")
	parser.add_option("--lvalert-listener-program", action = "append", default = [], metavar = "program", help = "set the programs to respond to lvalerts from this analysis, can be given multiple times")
	parser.add_option("--coincidence-threshold", metavar = "value", type = "float", default = 0.005, help = "Set the coincidence window in seconds (default = 0.005).  The light-travel time between instruments will be added automatically in the coincidence test.")
	parser.add_option("--likelihood-snapshot-interval", type = "float", metavar = "seconds", help = "How often to reread the marginalized likelihoood data and snapshot the trigger files.")

	options, filenames = parser.parse_args()

	fail = ""
	for option in ("bank_cache", "gracedb_far_threshold"):
		if getattr(options, option) is None:
			fail += "must provide option %s\n" % (option)
	if fail: raise ValueError, fail

	if options.copy_likelihoods and options.likelihood_file is None:
		raise ValueError("Must include --likelihood-file when giving --copy-likelihoods")

	#FIXME add consistency check?
	bankcache = inspiral_pipe.parse_cache_str(options.bank_cache)
	channel_dict = datasource.channel_dict_from_channel_list(options.channel)

	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)
	
	return options, filenames, bankcache, channel_dict


#
# MAIN
#


options, filenames, bank_cache, channel_dict = parse_command_line()

try: os.mkdir("logs")
except: pass
dag = dagparts.CondorDAG("trigger_pipe")


#
# setup the job classes
#


# Figure out if it is iir or not
if options.do_iir_pipeline is not None:
	gstlalInspiralJob = inspiral_pipe.generic_job('gstlal_iir_inspiral', condor_commands = {"want_graceful_removal":"True", "kill_sig":"15", "+Online_CBC_SVD":"True", "Requirements":"(TARGET.Online_CBC_SVD =?= True)"})
else:
	gstlalInspiralJob = inspiral_pipe.generic_job('gstlal_inspiral', condor_commands = {"want_graceful_removal":"True", "kill_sig":"15", "+Online_CBC_SVD":"True", "Requirements":"(TARGET.Online_CBC_SVD =?= True)"})
# A local universe job that will run in a loop marginalizing all of the likelihoods
margJob = inspiral_pipe.generic_job('gstlal_inspiral_marginalize_likelihoods_online')
# an lvalert_listen job
listenJob = lvalert_listen_job("lvalert_listen", gracedb_group = options.gracedb_group, gracedb_type = options.gracedb_type, progs = options.lvalert_listener_program)
# get urls job 
urlsJob = inspiral_pipe.generic_job("gstlal_ll_inspiral_get_urls")


#
# loop over banks to run gstlal inspiral pre clustering and far computation
#


listenNode = lvalert_listen_node(listenJob, dag)

for num_insp_nodes, svd_banks in enumerate(inspiral_pipe.build_bank_groups(bank_cache, [1], options.max_jobs)):

	# make a new likelihood file
	likefile = os.path.split(options.likelihood_file)[1]
	path = os.getcwd()
	likefile = "%s/%04d_%s" % (path, num_insp_nodes, likefile)
	if options.copy_likelihoods:
		shutil.copyfile(options.likelihood_file, likefile)
	svd_bank_string = ",".join([":".join([k, v[0]]) for k,v in svd_banks.items()])

	#FIXME add support for injections too
	inspNode = inspiral_pipe.generic_node(gstlalInspiralJob, dag, [],
		opts = {"psd-fft-length":options.psd_fft_length,
			"ht-gate-threshold":options.ht_gate_threshold,
			"channel-name":datasource.pipeline_channel_list_from_channel_dict(channel_dict),
			"state-vector-on-bits":datasource.state_vector_on_off_list_from_bits_dict(options.state_vector_on_off_dict),
			"svd-bank":svd_bank_string,
			"tmp-space":inspiral_pipe.condor_scratch_space(),
			"track-psd":"",
			"control-peak-time":options.control_peak_time,
			"coincidence-threshold":options.coincidence_threshold,
			"fir-stride":options.fir_stride,
			"data-source":options.data_source,
			"gracedb-far-threshold":options.gracedb_far_threshold,
			"gracedb-group":options.gracedb_group,
			"gracedb-type":options.gracedb_type,
			"thinca-interval":options.thinca_interval,
			"job-tag":"%04d" % num_insp_nodes,
			"likelihood_snapshot_interval":options.likelihood_snapshot_interval
			},
		input_files = {"marginalized-likelihood-file":options.marginalized_likelihood_file},
		output_files = {"output":"not_used.xml.gz",
				"likelihood-file":likefile
			}
		)
	print likefile

urlsNode = inspiral_pipe.generic_node(urlsJob, dag, [], opts = {"":". 0001,%04d 25" % (num_insp_nodes+1)}, input_files = {}, output_files = {})
margNode = inspiral_pipe.generic_node(margJob, dag, [], opts = {"":os.getcwd()}, input_files = {"":options.marginalized_likelihood_file}, output_files = {})


#
# Write out the dag and other flies
#


dag.write_sub_files()
# we probably want these jobs to retry indefinitely on dedicated nodes. A user
# can intervene and fix a problem without having to bring the dag down and up.
# There are few enough total jobs that this really shouldn't bog down the
# scheduler. For now 10000 will be considered indefinite
[node.set_retry(10000) for node in dag.get_nodes()]
dag.write_dag()
dag.write_script()
dag.write_cache()


#
# set up the webpage cgi scripts
# FIXME don't hardcode this stuff
#


shutil.copy2(dagparts.which('gstlal_llcbcsummary'), os.path.expanduser("~/public_html/cgi-bin"))
shutil.copy2(dagparts.which('gstlal_llcbcnode'), os.path.expanduser("~/public_html/cgi-bin"))
print >>sys.stderr, "\n\n NOTE! You can monitor the analysis at this url: https://ldas-jobs.ligo.caltech.edu/~%s/cgi-bin/gstlal_llcbcsummary?id=0001,%04d&dir=%s \n\n" % (os.environ['USER'], num_insp_nodes+1, os.getcwd())
