#!/usr/bin/env python
"""Make segment file to blind the results from foreground related triggers """

import os, argparse, logging, pycbc.version, h5py
from six.moves.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')
