#!/usr/bin/python3.11
"""Make segment file to blind the results from foreground related triggers """

import os, argparse, logging, pycbc.version, h5py
from urllib.parse import urlunparse
import pycbc.events
from pycbc.workflow import SegFile

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--version', action='version', version=pycbc.version.git_verbose_msg)
parser.add_argument('--verbose', action='store_true')
parser.add_argument('--foreground-triggers',
                    help="HDF file containing the zerolag foreground triggers "
                         "from the analysis")
parser.add_argument('--veto-file',
                    help="Baseline veto information that is added to the outptut")
parser.add_argument('--segment-name',
                    help="Segment name to use from the input veto file")
parser.add_argument('--output-file', help='Name of the output segment file')
parser.add_argument('--output-segment-name',
                    help="(optional), Name of output segment file list",
                    default="censor_foreground")
args = parser.parse_args()

pycbc.init_logging(args.verbose)

logging.info('Start')

f = h5py.File(args.foreground_triggers, 'r')

start = f['segments/foreground_veto/start'][:]
end = f['segments/foreground_veto/end'][:]
vsegs = pycbc.events.start_end_to_segments(start, end)

logging.info('Read in foreground veto segments')

# 2-ifo old style format
if 'detector_1' in f.attrs:
    ifo1, ifo2 = f.attrs['detector_1'], f.attrs['detector_2']
    ifos = [ifo1, ifo2]
# Multi-ifo format file
else:
    ifos = f.attrs['ifos'].split(' ')

fsegs, names = [], []
for ifo in ifos:
    segs = pycbc.events.select_segments_by_definer(args.veto_file, args.segment_name, ifo)
    logging.info('Read in veto segments from %s' % ifo)
    fsegs += [segs.coalesce() + vsegs.coalesce()]
    names += [args.output_segment_name]

file_url = urlunparse(['file', 'localhost',
                       os.path.abspath(args.output_file), None, None, None])
SegFile.from_multi_segment_list('UNUSED', fsegs, names, ifos, file_url=file_url)
logging.info('Done')
