#!/usr/bin/env python
#
# Copyright (C) 2010  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.

"""
This program makes a dag to generate svd banks
"""

__author__ = 'Chad Hanna <chad.hanna@ligo.org>'

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

import sys, os, copy, math
import subprocess, socket, tempfile

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

from glue import iterutils
from glue import pipeline
from glue import lal
from glue.ligolw import lsctables
from glue import segments
from optparse import OptionParser
from gstlal import inspiral_pipe
from gstlal import far 

## @file gstlal_inspiral_svd_bank_pipe
# This program will make a HTCondor DAG to automate the creation of svd bank files; see gstlal_inspiral_svd_bank_pipe for more information
#
#
# ### Graph of the condor DAG
#
# @dot
# digraph G {
#	// graph properties
#
#	rankdir=LR;
#	compound=true;
#	node [shape=record fontsize=10 fontname="Verdana"];     
#	edge [fontsize=8 fontname="Verdana"];
#
#	// nodes
#
#	"gstlal_svd_bank" [URL="\ref gstlal_svd_bank"];
# }
# @enddot
#
# This DAG implements only a single job type; gstlal_svd_bank
#
# ### Usage cases
#
# - Typical usage case for H1
#
#		$ gstlal_inspiral_svd_bank_pipe --autocorrelation-length 351 --instrument H1 --reference-psd reference_psd.xml --bank-cache H1_split_bank.cache --overlap 10 --flow 15 --num-banks 1 --output-name H1_bank
#
# - Please add more!
#
# ### Command line options
# 
#	+ `--instrument` [ifo]: set the name of the instrument, required
#	+ `--reference-psd` [file]: Set the name of the reference psd file, required
#	+ `--bank-cache` [file]: Set the name of the bank cache, required
#	+ `--overlap` [int]: Set the factor that describes the overlap of the sub banks, must be even!
#	+ `--identity-transform`: Turn off the SVD and use the identity reconstruction matrix
#	+ `--autocorrelation-length` [int]: The number of samples to use for auto-chisquared, default 201 should be odd
#	+ `--samples-min` [int]: The minimum number of samples to use for time slices default 1024
#	+ `--samples-max-256` [int]: The maximum number of samples to use for time slices with frequencies above 256Hz, default 1024
#	+ `--samples-max-64` [int]: The maximum number of samples to use for time slices with frequencies between 64Hz and 256 Hz, default 2048
#	+ `--samples-max` [int]: The maximum number of samples to use for time slices with frequencies below 64Hz, default 4096
#	+ `--tolerance` [float]: Set the SVD tolerance, default 0.9995
#	+ `--flow` [float]: Set the low frequency cutoff, default 40 (Hz)
#	+ `--num-banks [str]: The number of banks per job. can be given as a list like 1,2,3,4 then it will split up the bank cache into N groups with M banks each.")
#	+ `--output-name` [file]: Set the base name of the output, required
#	+ `--verbose`: Be verbose.
#
# ### Review Status
#
# | Reviewers				                | Hash				                       | Date 		| Diff to Head of Master     |
# | --------------------------------------- | ---------------------------------------- | ---------- | -------- |
# | Florent, Duncan Me., Jolien, Kipp, Chad | 67f86a6f2830f399c8d7c4cec6f357940a4b0abb | 2014-04-29	| <a href="@gstlal_inspiral_cgit_diff/bin/gstlal_inspiral_svd_bank_pipe?id=HEAD&id2=67f86a6f2830f399c8d7c4cec6f357940a4b0abb">gstlal_inspiral_svd_bank_pipe</a> |

def parse_command_line():
	parser = OptionParser()
	parser.add_option("--instrument", help = "set the name of the instrument, required")
	parser.add_option("--reference-psd", metavar = "file", help = "Set the name of the reference psd file, required")
	parser.add_option("--bank-cache", metavar = "file", help = "Set the name of the bank cache, required")
	parser.add_option("--overlap", metavar = "num", type = "int", default = 0, help = "set the factor that describes the overlap of the sub banks, must be even!")
	parser.add_option("--identity-transform", default = False, action = "store_true", help = "turn off the SVD and use the identity reconstruction matrix")
	parser.add_option("--autocorrelation-length", type = "int", default = 201, help = "The number of samples to use for auto-chisquared, default 201 should be odd")
	parser.add_option("--samples-min", type = "int", default = 1024, help = "The minimum number of samples to use for time slices default 1024")
	parser.add_option("--samples-max-256", type = "int", default = 1024, help = "The maximum number of samples to use for time slices with frequencies above 256Hz, default 1024")
	parser.add_option("--samples-max-64", type = "int", default = 2048, help = "The maximum number of samples to use for time slices with frequencies above 64Hz, default 2048")
	parser.add_option("--samples-max", type = "int", default = 4096, help = "The maximum number of samples to use for time slices with frequencies below 64Hz, default 4096")
	parser.add_option("--tolerance", metavar = "float", type = "float", default = 0.9995, help = "set the SVD tolerance, default 0.9995")
	parser.add_option("--flow", metavar = "num", type = "float", default = 40, help = "set the low frequency cutoff, default 40 (Hz)")
	parser.add_option("--output-name", help = "set the base name of the output, required")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	parser.add_option("--num-banks", metavar = "str", help = "the number of banks per job. can be given as a list like 1,2,3,4 then it will split up the bank cache into N groups with M banks each.")
	options, filenames = parser.parse_args()

	if options.overlap % 2:
		raise ValueError("overlap must be even")

	options.num_banks = [int(s) for s in options.num_banks.split(",")]

	return options, filenames


#
# get input arguments
#


options, filenames = parse_command_line()
ifo = options.instrument
ref_psd = options.reference_psd


#
# Do some setup
#


try:
	os.mkdir("logs")
except:
	pass
dag = inspiral_pipe.DAG(options.output_name)
svdJob = inspiral_pipe.generic_job("gstlal_svd_bank", tag_base = "gstlal_svd_bank_%s" % ifo, condor_commands = {"request_memory":"1999"})
# Assumes cache is sorted by chirpmass or whatever the SVD sorting algorithm that was chosen
files = [lal.CacheEntry(line).path for line in open(options.bank_cache)]


#
# loop over files to set up svd bank jobs
#

groups = list(inspiral_pipe.group(files, options.num_banks))
bank_ids = [0]
for i, f in enumerate(groups):
	# handle the edges by not clipping so you retain the template bank as intended.  
	clipleft = [options.overlap / 2] * len(f) # overlap must be even
	clipright = [options.overlap / 2] * len(f) # overlap must be even
	if i == 0:
		clipleft[0] = 0
	if i == len(groups) - 1:
		clipright[-1] = 0
	bank_ids = range(bank_ids[-1] + 1, bank_ids[-1] + 1 + len(f))
	svd_bank_name = inspiral_pipe.T050017_filename(ifo, "GSTLAL_SVD_BANK_%d" % i, 0, 0, ".xml.gz", path = svdJob.output_path)
	svd_bank_name = os.path.join(os.getcwd(), svd_bank_name)
	dag.output_cache.append(lal.CacheEntry(ifo, "GSTLAL_SVD_BANK_%d" % i, segments.segment(0, 0), "file://localhost%s" % (svd_bank_name,)))
	
	svdNode = inspiral_pipe.generic_node(svdJob, dag, [],
		opts = {"snr-threshold":far.ThincaCoincParamsDistributions.snr_min,
			"flow":options.flow,
			"svd-tolerance":options.tolerance,
			"ortho-gate-fap":0.5,
			"samples-min":options.samples_min,
			"samples-max":options.samples_max,
			"samples-max-64":options.samples_max_64,
			"samples-max-256":options.samples_max_256,
			"clipleft":clipleft,
			"clipright":clipright,
			"autocorrelation-length":options.autocorrelation_length,
			"bank-id":bank_ids
			},
		input_files = {
			"template-bank":f,
			"reference-psd":ref_psd
			},
		output_files = {
			"write-svd-bank":svd_bank_name
			}
		)
	if options.identity_transform:
		svdNode.add_var_arg("--identity-transform")


#
# Write out the dag files
#


dag.write_sub_files()
dag.write_dag()
dag.write_script()
dag.write_cache()
