#!/usr/bin/python3.13
#
# Copyright (C) 2020-2021  Kipp Cannon, Patrick Godwin, Chad Hanna, 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_segments

#
# Define/parse command line options
#

parser = argparse.ArgumentParser(description='Download GWOSC 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("-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("-f", "--flag", metavar="flag", action="append", default=[], help=f"Set the flags to use to query segments, one for each instrument.")
parser.add_argument("-o", "--output", metavar="file", default="segments.xml.gz", help="Set the name of the segments file generated. Default: segments.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))])

flags = {}
for flag in args.flag:
	ifo, _ = flag.split(":", 1)
	flags[ifo] = flag

#
# Query segments
#

print("Downloading segment lists ...", file=sys.stderr)
segments = query_dqsegdb_segments(instruments, args.start, args.end, flags=flags, server=args.server)

#
# Write segments 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(segments, "datasegments")
process.set_end_time_now()

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