#!/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')
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]

# FIXME: Don't use URL once Larne supports taking file paths directly in the
#        workflow File class
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')