#!/usr/bin/python3.13
#
# Copyright (C) 2020-2021  Heather Fong, Patrick Godwin, Ryan Magee, Cody Messick
#
# 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 argparse
import sys

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

from gstlal.segments import DEFAULT_DQSEGDB_SERVER, query_dqsegdb_veto_segments

#
# Define/parse command line options
#

parser = argparse.ArgumentParser(description='Download GWOSC veto segments.')
parser.add_argument("start", metavar="START", type=int, help="Set the start GPS time.")
parser.add_argument("end", metavar="END", type=int, help="Set the end GPS time.")
parser.add_argument("ifo", metavar="IFO", nargs="+", help="Set the instruments to process, e.g. H1.")
parser.add_argument("-v", "--veto-definer-file", metavar="file", help="The veto definer file in which to query DQSegDB for.")
parser.add_argument("-s", "--server", metavar="url", default=DEFAULT_DQSEGDB_SERVER, help=f"Set the URL of the DQSegDB server to query. Default: {DEFAULT_DQSEGDB_SERVER}")
parser.add_argument("-c", "--category", help="Set the veto category to generate veto segments for.")
parser.add_argument("--cumulative", action="store_true", help="Set whether veto categories are cumulative, e.g. if CAT2 and cumulative, combine both CAT1 and CAT2 vetoes.")
parser.add_argument("-o", "--output", metavar="file", default="vetoes.xml.gz", help="Set the name of the segments file generated. Default: vetoes.xml.gz")
args = parser.parse_args()

instruments = set([])
for x in args.ifo:
	# assume all instruments are two characters and split combos like H1L1 if they exist
	instruments |= set([x[i*2:i*2+2] for i in range(int(len(x) / 2))])

#
# Query veto segments
#

print("Downloading veto segments ...", file=sys.stderr)
vetoes = query_dqsegdb_veto_segments(
	instruments,
	args.start,
	args.end,
	args.veto_definer_file,
	category=args.category,
	cumulative=args.cumulative,
	server=args.server,
)

#
# Write vetoes to disk
#

xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())
process = ligolw_process.register_to_xmldoc(xmldoc, sys.argv[0], vars(args))
with ligolw_segments.LigolwSegments(xmldoc, process) as lwseglists:
	lwseglists.insert_from_segmentlistdict(vetoes, "vetoes")
process.set_end_time_now()

ligolw_utils.write_filename(xmldoc, args.output, verbose=True)
xmldoc.unlink()
