#!/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('--strict-coinc-time', action='store_true',
                    help="Veto any time that is vetoed in either detector")
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')

# Check if using input files from old (2-ifo) coinc code
if 'segments/foreground_veto/start' in f:
    start = f['segments/foreground_veto/start'][:]
    end = f['segments/foreground_veto/end'][:]
    vsegs = pycbc.events.start_end_to_segments(start, end)

    ifo1, ifo2 = f.attrs['detector_1'], f.attrs['detector_2']

    ifos, fsegs, names = [], [], []
    for ifo in [ifo1, ifo2]:
        segs = pycbc.events.select_segments_by_definer(args.veto_file, args.segment_name, ifo)
        fsegs += [(segs.coalesce() + vsegs.coalesce()).coalesce()]
        names += [args.output_segment_name]
        ifos += [ifo]

    if args.strict_coinc_time:
        fsegs[0] = (fsegs[0].coalesce() + fsegs[1].coalesce()).coalesce()
        fsegs[1] = fsegs[0]

else:
    # Version used for multi-ifo coinc code
    if args.strict_coinc_time:
        raise RuntimeError("Can't use strict coinc time option with multi-ifo pipeline!")
    ifolist = f.attrs['ifos'].split(' ')
    ifokey = ''.join(sorted(ifolist))

    start = f['segments/%s/foreground_veto/start' % ifokey][:]
    end = f['segments/%s/foreground_veto/end' % ifokey][:]
    mvsegs = pycbc.events.start_end_to_segments(start, end)

    for key in f['segments']:
        start = f['segments/%s/foreground_veto/start' % key][:]
        end = f['segments/%s/foreground_veto/end' % key][:]
        vsegs = pycbc.events.start_end_to_segments(start, end)
        mvsegs = (mvsegs & vsegs).coalesce()

    ifos, fsegs, names = [], [], []
    for ifo in ifolist:
        segs = pycbc.events.select_segments_by_definer(args.veto_file, args.segment_name, ifo)
        fsegs += [(segs.coalesce() + mvsegs.coalesce()).coalesce()]
        names += [args.output_segment_name]
        ifos += [ifo]

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')
