#!/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) 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.

from optparse import OptionParser
import sys

from ligo.lw import ligolw
from ligo.lw import lsctables
from ligo.lw import utils as ligolw_utils
from ligo.lw.utils import segments as ligolw_segments
from ligo.lw.utils import process as ligolw_process
from ligo.segments import utils as segmentsUtils

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

def parse_command_line():
	parser = OptionParser()
	parser.add_option("--segment-name", metavar = "name", action = "append", help = "Set the name of segments default datasegments (can be given multiple times)")
	parser.add_option("--intersection", action = "store_true", help = "intersection of --segment-name across provided segment files")
	parser.add_option("--union", action = "store_true", help = "union of --segment-name across provided segment files")
	parser.add_option("--diff", action = "store_true", help = "diff of --segment-name across provided segment files (cannot provide more than 2 segment files if invoking --diff)")
	parser.add_option("--output-file", metavar = "file", help = "Set the name of the output file")
	parser.add_option("--output-segment-name", metavar = "name", action = "append", help = "Set the name of output segments default datasegments (can be given multiple times)")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	options, filenames = parser.parse_args()
	
	if len([op for op in (options.intersection, options.union, options.diff) if op is not None]) != 1:
		raise ValueError("must specify exactly one of --intersection, --union or --diff")

	for option in ("segment_name", "output_segment_name"):
		if not getattr(options, option):
			setattr(options, option, ["datasegments"])

	assert len(options.output_segment_name) == len(options.segment_name), "must provide same number of input and output segment names"

	if options.diff:
		assert len(filenames) == 2, "must provide exactly 2 segment files for diff operation"
	else:
		assert len(filenames) >= 2, "must provide at least 2 segment files, %d provided" % len(filenames)

	return options, filenames

options, filenames = parse_command_line()

def extract_segs(fname, segnames):
	llwsegments = ligolw_segments.LigolwSegments(ligolw_utils.load_filename(fname, verbose = True, contenthandler = LIGOLWContentHandler))
	return dict((segname, llwsegments.get_by_name(segname).coalesce()) for segname in segnames)

segsdicts = {}
for segfile in filenames:
	segsdicts[segfile] = extract_segs(segfile, options.segment_name)

if options.verbose:
	for segfile in filenames:
		print("\n%s" % segfile)
		for segname, segslist in segsdicts[segfile].items():
			print("%s live time: %s, two or more detectors on: %f" % (segname, abs(segslist), abs(segmentsUtils.vote(segslist.values(), 2))))

segslist = []
for segname in options.segment_name:
	segslist.append(segsdicts[filenames[0]][segname])
	for segfile in filenames[1:]:
		if options.union:
			segslist[-1] |= segsdicts[segfile][segname]
		elif options.intersection:
			segslist[-1] &= segsdicts[segfile][segname]
		elif options.diff:
			segslist[-1] -= segsdicts[segfile][segname]

if options.verbose:
	print("\nresults")
	for i, segs in enumerate(segslist):
		print("%s live time: %s, two or more detectors on: %f" % (options.output_segment_name[i], abs(segs), abs(segmentsUtils.vote(segs.values(), 2))))

if options.output_file:
	xmldoc = ligolw.Document()
	xmldoc.appendChild(ligolw.LIGO_LW())

	process = ligolw_process.register_to_xmldoc(xmldoc, sys.argv[0], options.__dict__)

	lwseglists = ligolw_segments.LigolwSegments(xmldoc)
	for i, segs in enumerate(segslist):
		lwseglists.insert_from_segmentlistdict(segs, options.output_segment_name[i])
	lwseglists.optimize()
	lwseglists.finalize(process)

	process.set_end_time_now()
	ligolw_utils.write_filename(xmldoc, options.output_file, verbose = options.verbose)

