#!/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 sys, os
from ligo.lw import ligolw
from ligo.lw import utils as ligolw_utils
from ligo.lw import lsctables
import numpy
import argparse
import h5py
from gstlal import pipeline
from gstlal import dagparts

class LIGOLWContentHandler(ligolw.LIGOLWContentHandler):
	pass
lsctables.use_in(LIGOLWContentHandler)

def file_len(fname):
      with open(fname) as f:
            for i, l in enumerate(f):
                  pass
      return i+1

parser = argparse.ArgumentParser()
parser.add_argument("--psd-xml-file", help = "provide a psd xml file")
parser.add_argument("--bank-file", help = "provide the bank file for which overlaps will be calculated")
parser.add_argument("--split-bank-cache", help = "Cache file containing paths to split banks.")
args = parser.parse_args()

xmldoc = ligolw_utils.load_filename(args.bank_file, verbose = True, contenthandler = LIGOLWContentHandler)
sngl_inspiral_table = lsctables.SnglInspiralTable.get_table(xmldoc)

number_of_templates = len(sngl_inspiral_table)

try:
	os.mkdir("logs")
except:
	pass
dag = dagparts.DAG("metric_overlap")

overlapJob = dagparts.DAGJob("gstlal_inspiral_metric_overlap", condor_commands = {"want_graceful_removal":"True", "kill_sig":"15", "accounting_group":"ligo.prod.o3.cbc.uber.gstlaloffline"})
addJob = dagparts.DAGJob("gstlal_inspiral_add_metric_overlaps", condor_commands = {"want_graceful_removal":"True", "kill_sig":"15", "accounting_group":"ligo.prod.o3.cbc.uber.gstlaloffline"})

overlapnodes = []

for i in range(0, file_len(args.split_bank_cache)):
      overlapnodes.append(dagparts.DAGNode(overlapJob, dag, parent_nodes = [], opts = {"split-bank-cache":args.split_bank_cache, "split-bank-index": i, "number-of-templates": number_of_templates}, output_files = {"out-h5-file":"%s/metric_overlaps_%d.h5" % (overlapJob.output_path, i)}, input_files = {"psd-xml-file": args.psd_xml_file, "bank-file": args.bank_file}))

addnode = dagparts.DAGNode(addJob, dag, parent_nodes = overlapnodes, output_files = {"out-h5-file": "overlaps.h5"}, input_files = {"": [n.output_files["out-h5-file"] for n in overlapnodes]})

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