#!/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
# Copyright (C) 2015  Surabhi Sachdev, Tjonnie Li
#
# 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.

###############################################################################
#
# IMPORT MODULES
#
###############################################################################

from optparse import OptionParser
import os

from lal.utils import CacheEntry
from ligo.lw import lsctables
from ligo.lw import utils
from ligo.lw import ligolw
from ligo.lw.utils import process as ligolw_process
from ligo.segments import segment
import numpy

from gstlal.datafind import DataType

###############################################################################
#
# IMPORT MODULES
#
###############################################################################

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

###############################################################################
#
# COMMAND LINE PARSING
#
###############################################################################

def parse_command_line():
	parser = OptionParser()
	parser.add_option("-o", "--output-path", metavar = "path", default = ".", help = "Set the path to the directory where output files will be written.  Default is \".\".")
	parser.add_option("-u", "--usertag", metavar = "usertag", help = "Set the user tag",default="INJSPLITTER")
	parser.add_option("-n", "--nsplit", metavar = "count", type = "int", help = "Number of files you want the original file to be split into",default=1)
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	options, filenames = parser.parse_args()

	if len(filenames)!=1:
		raise ValueError("Provide one injection file")

	return options, filenames

options, filenames = parse_command_line()

###############################################################################
#
# MAIN CODE
#
###############################################################################

# GETTING COMMAND LINE OPTIONS FOR PRINTING INTO THE TABLE
opts_dict = dict((k, v) for k, v in options.__dict__.items() if v is not False and v is not None)

# LOAD INJECTION TABLE
xmldoc=utils.load_filename(filenames[0], verbose = options.verbose, contenthandler=ContentHandler)
sim_inspiral_table=lsctables.SimInspiralTable.get_table(xmldoc)
# Sort the sim inspiral table based on geocent end time
sim_inspiral_table.sort(key = lambda row: (row.geocent_end_time + 1e-9 * row.geocent_end_time_ns))
process_params_table = lsctables.ProcessParamsTable.get_table(xmldoc)

# create directory if needed
os.makedirs(options.output_path, exist_ok=True)

# infer metadata from injection filename
try:
	entry = CacheEntry.from_T050017(filenames[0])
except ValueError:
	ifos = "H1K1L1V1"
	span = segment(0, 0)
else:
	ifos = entry.observatory
	span = entry.segment

# PREPARE PROCESS TABLE WITH INFORMATION ABOUT THE CURRENT PROGRAM
xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())
node = xmldoc.childNodes[-1]
process = ligolw_process.register_to_xmldoc(xmldoc,
"gstlal_injsplitter", opts_dict,
version="no version", cvs_repository="gstlal")
node.appendChild(process_params_table)

sim_inspiral_table_split = lsctables.New(lsctables.SimInspiralTable)
node.appendChild(sim_inspiral_table_split)

evensplit = numpy.array_split(sim_inspiral_table,options.nsplit);
for i in range(options.nsplit):
	sim_inspiral_table_split[:] = evensplit[i]
	process.set_end_time_now()
	inj_filename = DataType.SPLIT_INJECTIONS.filename(ifos, span, f"{i:04d}", options.usertag)
	utils.write_filename(xmldoc, os.path.join(options.output_path, inj_filename), verbose = options.verbose)
