#!/usr/bin/python3.11
"""
The program combines coincident output files generated
by pycbc_coinc_findtrigs to generated a mapping between SNR and FAP, along
with producing the combined foreground and background triggers
"""
import argparse, h5py, logging, itertools, copy, pycbc.io, numpy, lal
from pycbc.events import veto, coinc, significance
import pycbc.version
import pycbc.conversions as conv

parser = argparse.ArgumentParser()
# General required options
parser.add_argument('--verbose', action='count')
parser.add_argument('--version', action='version',
                    version=pycbc.version.git_verbose_msg)
parser.add_argument('--cluster-window', type=float, default=10,
                    help='Length of time window in seconds to cluster coinc '
                         'events [default=10s]')
parser.add_argument('--zero-lag-coincs', nargs='+',
                    help='Files containing the injection zerolag coincidences')
parser.add_argument('--full-data-background',
                    help='background file from full data for use in analyzing '
                         'injection coincs')
parser.add_argument('--veto-window', type=float, default=.1,
                    help='Time around each zerolag trigger to window out '
                         '[default=.1s]')
parser.add_argument('--ifos', nargs='+',
                    help='List of ifos used in these coincidence files')
significance.insert_significance_option_group(parser)
parser.add_argument('--output-file')
args = parser.parse_args()

significance.check_significance_options(args, parser)

if args.verbose:
    log_level = logging.INFO
    logging.basicConfig(format='%(asctime)s : %(message)s', level=log_level)

ifo_key = ''.join(args.ifos)
significance_dict = significance.digest_significance_options([ifo_key], args)

window = args.cluster_window
logging.info("Loading coinc zerolag triggers")
zdata = pycbc.io.MultiifoStatmapData(files=args.zero_lag_coincs, ifos=args.ifos)
zdata = zdata.cluster(window)

f = h5py.File(args.output_file, "w")

f.attrs['num_of_ifos'] = zdata.attrs['num_of_ifos']
f.attrs['pivot'] = zdata.attrs['pivot']
f.attrs['fixed'] = zdata.attrs['fixed']
f.attrs['timeslide_interval'] = zdata.attrs['timeslide_interval']
f.attrs['ifos'] = ' '.join(sorted(args.ifos))

# Copy over the segment for coincs and singles
for key in zdata.seg.keys():
    f['segments/%s/start' % key] = zdata.seg[key]['start'][:]
    f['segments/%s/end' % key] = zdata.seg[key]['end'][:]

logging.info('writing zero lag triggers')
if len(zdata) > 0:
    for key in zdata.data:
        f['foreground/%s' % key] = zdata.data[key]
else:
    for key in zdata.data:
        f['foreground/%s' % key] = numpy.array([], dtype=zdata.data[key].dtype)

logging.info('calculating statistics excluding zerolag')
fb = h5py.File(args.full_data_background, "r")

# we expect the injfull file to contain injection data as pivot
# and fullinj to contain full data as pivot
background_time = float(fb.attrs['background_time'])
coinc_time = float(fb.attrs['foreground_time'])
back_stat = fb['background_exc/stat'][:]
dec_fac = fb['background_exc/decimation_factor'][:]

f.attrs['background_time_exc'] = background_time
f.attrs['foreground_time_exc'] = coinc_time
f.attrs['background_time'] = background_time
f.attrs['foreground_time'] = coinc_time

if len(zdata) > 0:

    _, fg_far_exc = significance.get_far(
        back_stat,
        zdata.stat,
        dec_fac,
        background_time,
        **significance_dict[ifo_key])

    fg_far_exc = significance.apply_far_limit(fg_far_exc, significance_dict, combo=ifo_key)

    ifar_exc = 1. / fg_far_exc
    fap_exc = 1 - numpy.exp(- coinc_time / ifar_exc)
    f['foreground/ifar_exc'] = conv.sec_to_year(ifar_exc)
    f['foreground/fap_exc'] = fap_exc

else:
    f['foreground/ifar_exc'] = numpy.array([])
    f['foreground/fap_exc'] = numpy.array([])

logging.info("Done")
