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

import os, argparse, logging, pycbc.version, h5py, urlparse
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 = urlparse.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')
