#!/home/conda/feedstock_root/build_artifacts/bld/rattler-build_gstlal-ugly_1767612080/host_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_pl/bin/python
import itertools
import numpy, math
from ligo.lw import ligolw
from ligo.lw import lsctables
from ligo.lw import table
from ligo.lw.utils import process as ligolw_process
from gstlal import metric as metric_module
import os,sys,argparse
from gstlal import inspiral
from gstlal import dagparts

# Read command line options
def parse_command_line():

	parser = argparse.ArgumentParser(description="Template generator via tree.")

	# Options for generating the template bank
	parser.add_argument("--min-mass1", action="store", type=float,\
				    default=3.0, help="Minimum mass1 to generate bank.")
	parser.add_argument("--max-mass1", action="store", type=float,\
				    default=10.0, help="Maximum mass1 to generate bank.")
	parser.add_argument("--min-mass2", action="store", type=float,\
				    default=3.0, help="Minimum mass2 to generate bank.")
	parser.add_argument("--max-mass2", action="store", type=float,\
				    default=10.0, help="Maximum mass2 to generate bank.")
	parser.add_argument("--min-spin1x", action="store", type=float,\
				    default=0, help="Minimum mass1 to generate bank.")
	parser.add_argument("--max-spin1x", action="store", type=float,\
				    default=0, help="Maximum mass1 to generate bank.")
	parser.add_argument("--min-spin2x", action="store", type=float,\
				    default=0, help="Minimum mass2 to generate bank.")
	parser.add_argument("--max-spin2x", action="store", type=float,\
				    default=0, help="Maximum mass2 to generate bank.")
	parser.add_argument("--min-spin1y", action="store", type=float,\
				    default=0, help="Minimum mass1 to generate bank.")
	parser.add_argument("--max-spin1y", action="store", type=float,\
				    default=0, help="Maximum mass1 to generate bank.")
	parser.add_argument("--min-spin2y", action="store", type=float,\
				    default=0, help="Minimum mass2 to generate bank.")
	parser.add_argument("--max-spin2y", action="store", type=float,\
				    default=0, help="Maximum mass2 to generate bank.")
	parser.add_argument("--min-spin1z", action="store", type=float,\
			default=0, help="Minimum mass1 to generate bank.")
	parser.add_argument("--max-spin1z", action="store", type=float,\
				    default=0, help="Maximum mass1 to generate bank.")
	parser.add_argument("--min-spin2z", action="store", type=float,\
				    default=0, help="Minimum mass2 to generate bank.")
	parser.add_argument("--max-spin2z", action="store", type=float,\
				    default=0, help="Maximum mass2 to generate bank.")
	parser.add_argument("--min-match", action="store", type=float,\
				    default=0.95, help="Minimum match to generate bank.")
	parser.add_argument("--flow", action="store", type=float,\
				    default=30.0, help="Low frequency cutoff for overlap calculations.")
	parser.add_argument("--fhigh", action="store", type=float,\
				    default=1024.0, help="High frequency cutoff for overlap calculations.")
	parser.add_argument("--approximant", action="store", type=str,\
				    default="TaylorF2", help="Specify approximant.")
	parser.add_argument("--psd-file", action="store",\
				    default=None, help="Input PSD file.")
	parser.add_argument("--noise-model", action="store",\
				    default=None, help="Specify standard noise model.")
	parser.add_argument("--num-jobs", action="store",type=int,\
				    default=36, help="Specify the number of jobs to run. Default: 36. This is obeyed approximately")

	args = parser.parse_args()

	if args.noise_model and args.psd_file:
		raise ValueError("Cannot specify both --psd-file and --noise-model")

	if not (args.noise_model or args.psd_file):
		args.psd_file = "test.xml.gz" # FIXME for bw compatibility

	if args.noise_model:
		raise NotImplementedError("IMPLEMENT NOISE MODELS!!")

	return args

args = parse_command_line()
dag = dagparts.DAG("treebank")
treeJob = dagparts.DAGJob("gstlal_inspiral_treebank")
argsdict = dict((key.replace("_","-"), value) for key,value in vars(args).items())
mass2 = numpy.logspace(math.log(args.min_mass2, 8/3.), math.log(args.max_mass2, 8/3.), args.num_jobs, base=8/3.)
del argsdict["num-jobs"]
pnodes = []
cnt = 1
for minmass2,maxmass2 in zip(mass2[:-1],mass2[1:]):
	argsdict["min-mass2"] = minmass2
	argsdict["max-mass2"] = maxmass2
	argsdict["user-tag"] = cnt
	cnt+=1
	pnodes.append(dagparts.DAGNode(treeJob, dag, parent_nodes = [], input_files = {}, output_files = {}, opts = argsdict))

#FIXME add a ligolw_add job to the end of the dag
try:
	os.mkdir("logs")
except:
	pass

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

		




