#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Copyright (C) Joshua Smith (2016-)
#
# This file is part of the hveto python package.
#
# hveto is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# hveto is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with hveto.  If not, see <http://www.gnu.org/licenses/>.

"""Run the HierarchichalVeto (hveto) algorithm over some data
"""

import time
jobstart = time.time()  # noqa

import os
import sys
import json
import random
import warnings
import datetime
import configparser
import multiprocessing
from pathlib import Path
from socket import getfqdn
from getpass import getuser

import numpy

from matplotlib import use
use('agg')  # noqa

from gwpy.io.cache import read_cache
from gwpy.segments import (Segment, SegmentList,
                           DataQualityFlag, DataQualityDict)

from gwdetchar import cli
from gwdetchar.io.html import (FancyPlot, cis_link)
from gwdetchar.omega import batch
from gwdetchar.plot import texify

from hveto import (__version__, config, core, plot, html, utils)
from hveto.plot import (HEADER_CAPTION, ROUND_CAPTION, get_column_label)
from hveto.segments import (write_ascii as write_ascii_segments,
                            read_veto_definer_file)
from hveto.triggers import (get_triggers, find_auxiliary_channels)

IFO = os.getenv('IFO')

__author__ = 'Duncan Macleod <duncan.macleod@ligo.org>'
__credits__ = ('Joshua Smith <joshua.smith@ligo.org>, '
               'Alex Urban <alexander.urban@ligo.org>')

logger = cli.logger(name=os.path.basename(__file__))


# -- parse command line -------------------------------------------------------

def abs_path(p):
    return os.path.abspath(os.path.expanduser(p))


parser = cli.create_parser(description=__doc__, version=__version__)
cli.add_gps_start_stop_arguments(parser)
cli.add_ifo_option(parser, required=IFO is None, ifo=IFO)
parser.add_argument('-f', '--config-file', action='append', default=[],
                    type=abs_path,
                    help='path to hveto configuration file, can be given '
                         'multiple times (files read in order)')
cli.add_nproc_option(parser, default=1)
parser.add_argument('-p', '--primary-cache', default=None, type=abs_path,
                    help='path for cache containing primary channel files')
parser.add_argument('-a', '--auxiliary-cache', default=None, type=abs_path,
                    help='path for cache containing auxiliary channel files, '
                         'files contained must be T050017-compliant with the '
                         'channel name as the leading name parts, e.g. '
                         '\'L1-GDS_CALIB_STRAIN_<tag>-<start>-<duration>.'
                         '<ext>\' for L1:GDS-CALIB_STRAIN triggers')
parser.add_argument('-S', '--analysis-segments', action='append', default=[],
                    type=abs_path,
                    help='path to file containing segments for '
                         'the analysis flag (name in data file '
                         'must match analysis-flag in config file)')
parser.add_argument('-w', '--omega-scans', type=int, metavar='NSCAN',
                    help='generate a workflow of omega scans for each round, '
                         'this will launch automatically to condor, '
                         'requires the gwdetchar package')

pout = parser.add_argument_group('Output options')
pout.add_argument('-o', '--output-directory', default=os.curdir,
                  help='path of output directory, default: %(default)s')

args = parser.parse_args()

ifo = args.ifo
start = int(args.gpsstart)
end = int(args.gpsend)
duration = end - start

logger.info("-- Welcome to Hveto --")
logger.info("GPS start time: %d" % start)
logger.info("GPS end time: %d" % end)
logger.info("Interferometer: %s" % ifo)

# -- initialisation -----------------------------------------------------------

# read configuration
cp = config.HvetoConfigParser(ifo=ifo)
cp.read(args.config_file)
logger.info("Parsed configuration file(s)")

# format output directory
outdir = abs_path(args.output_directory)
if not os.path.isdir(outdir):
    os.makedirs(outdir)
os.chdir(outdir)
logger.info("Working directory: %s" % outdir)
segdir = 'segments'
plotdir = 'plots'
trigdir = 'triggers'
omegadir = 'scans'
for d in [segdir, plotdir, trigdir, omegadir]:
    if not os.path.isdir(d):
        os.makedirs(d)

# prepare html variables
htmlv = {
    'title': '%s Hveto | %d-%d' % (ifo, start, end),
    'config': None,
    'context': ifo.lower(),
}

# get segments
aflag = cp.get('segments', 'analysis-flag')
url = cp.get('segments', 'url')
padding = tuple(cp.getfloats('segments', 'padding'))
if args.analysis_segments:
    segs_ = DataQualityDict.read(args.analysis_segments, gpstype=float)
    analysis = segs_[aflag]
    span = SegmentList([Segment(start, end)])
    analysis.active &= span
    analysis.known &= span
    analysis.coalesce()
    logger.debug("Segments read from disk")
else:
    analysis = DataQualityFlag.query(aflag, start, end, url=url)
    logger.debug("Segments recovered from %s" % url)
if padding != (0, 0):
    mindur = padding[0] - padding[1]
    analysis.active = type(analysis.active)([s for s in analysis.active if
                                             abs(s) >= mindur])
    analysis.pad(*padding, inplace=True)
    logger.debug("Padding %s applied" % str(padding))
livetime = int(abs(analysis.active))
livetimepc = livetime / duration * 100.
logger.info("Retrieved %d segments for %s with %ss (%.2f%%) livetime"
            % (len(analysis.active), aflag, livetime, livetimepc))

# apply vetoes from veto-definer file
try:
    vetofile = cp.get('segments', 'veto-definer-file')
except configparser.NoOptionError:
    vetofile = None
else:
    try:
        categories = cp.getfloats('segments', 'veto-definer-categories')
    except configparser.NoOptionError:
        categories = None
    # read file
    vdf = read_veto_definer_file(vetofile, start=start, end=end, ifo=ifo)
    logger.debug("Read veto-definer file from %s" % vetofile)
    # get vetoes from segdb
    vdf.populate(source=url, segments=analysis.active, on_error='warn')
    # coalesce flags from chosen categories
    vetoes = DataQualityFlag('%s:VDF-VETOES:1' % ifo)
    nflags = 0
    for flag in vdf:
        if not categories or vdf[flag].category in categories:
            vetoes += vdf[flag]
            nflags += 1
    try:
        deadtime = int(abs(vetoes.active)) / int(abs(vetoes.known)) * 100
    except ZeroDivisionError:
        deadtime = 0
    logger.debug("Coalesced %ss (%.2f%%) of deadtime from %d veto flags"
                 % (abs(vetoes.active), deadtime, nflags))
    # apply to analysis segments
    analysis -= vetoes
    logger.debug("Applied vetoes from veto-definer file")
    livetime = int(abs(analysis.active))
    livetimepc = livetime / duration * 100.
    logger.info("%ss (%.2f%%) livetime remaining after vetoes"
                % (livetime, livetimepc))

snrs = cp.getfloats('hveto', 'snr-thresholds')
minsnr = min(snrs)
windows = cp.getfloats('hveto', 'time-windows')

# record all segments
segments = DataQualityDict()
segments[analysis.name] = analysis

# -- load channels ------------------------------------------------------------

# get primary channel name
pchannel = cp.get('primary', 'channel')

# read auxiliary cache
if args.auxiliary_cache is not None:
    acache = read_cache(args.auxiliary_cache)
else:
    acache = None

# load auxiliary channels
auxetg = cp.get('auxiliary', 'trigger-generator')
auxfreq = cp.getfloats('auxiliary', 'frequency-range')
try:
    auxchannels = cp.get('auxiliary', 'channels').strip('\n').split('\n')
except config.configparser.NoOptionError:
    auxchannels = find_auxiliary_channels(auxetg, (start, end), ifo=ifo,
                                          cache=acache)
    cp.set('auxiliary', 'channels', '\n'.join(auxchannels))
    logger.debug("Auto-discovered %d auxiliary channels" % len(auxchannels))
else:
    auxchannels = sorted(set(auxchannels))
    logger.debug("Read list of %d auxiliary channels" % len(auxchannels))


# load unsafe channels list
_unsafe = cp.get('safety', 'unsafe-channels')
if os.path.isfile(_unsafe):  # from file
    unsafe = set()
    with open(_unsafe, 'rb') as f:
        for c in f.read().rstrip('\n').split('\n'):
            if c.startswith('%(IFO)s'):
                unsafe.add(c.replace('%(IFO)s', ifo))
            elif not c.startswith('%s:' % ifo):
                unsafe.add('%s:%s' % (ifo, c))
            else:
                unsafe.add(c)
else:  # or from line-seprated list
    unsafe = set(_unsafe.strip('\n').split('\n'))
unsafe.add(pchannel)
cp.set('safety', 'unsafe-channels', '\n'.join(sorted(unsafe)))
logger.debug("Read list of %d unsafe channels" % len(unsafe))

# remove unsafe channels
nunsafe = 0
for i in range(len(auxchannels) - 1, -1, -1):
    if auxchannels[i] in unsafe:
        logger.warning("Auxiliary channel %r identified as unsafe and has "
                       "been removed" % auxchannels[i])
        auxchannels.pop(i)
        nunsafe += 1
logger.debug("%d auxiliary channels identified as unsafe" % nunsafe)
naux = len(auxchannels)
logger.info("Identified %d auxiliary channels to process" % naux)

# record INI file in output HTML directory
inifile = '%s-HVETO_CONFIGURATION-%d-%d.ini' % (ifo, start, duration)
if os.path.isfile(inifile) and any(
        os.path.samefile(inifile, x) for x in args.config_file):
    logger.debug("Cannot write INI file to %s, file was given as input")
else:
    with open(inifile, 'w') as f:
        cp.write(f)
    logger.info("Configuration recorded as %s" % inifile)
htmlv['config'] = inifile

# -- load primary triggers ----------------------------------------------------

# read primary cache
if args.primary_cache is not None:
    pcache = read_cache(args.primary_cache)
else:
    pcache = None

# load primary triggers
petg = cp.get('primary', 'trigger-generator')
psnr = cp.getfloat('primary', 'snr-threshold')
pfreq = cp.getfloats('primary', 'frequency-range')
preadkw = cp.getparams('primary', 'read-')
#if pcache is not None:  # auto-detect the file format # original line, but not None, we don't have auto-detection
if pcache is None:  # auto-detect the file format
    logger.debug('Unsetting the primary trigger file format')
    preadkw['format'] = None
    preadkw['path'] = 'triggers' # original line, YMKIM, 2020.05.06
ptrigfindkw = cp.getparams('primary', 'trigfind-')
primary = get_triggers(pchannel, petg, analysis.active, snr=psnr, frange=pfreq,
                       cache=pcache, nproc=args.nproc,
                       trigfind_kwargs=ptrigfindkw, **preadkw)
fcol, scol = primary.dtype.names[1:3]

if len(primary):
    logger.info("Read %d events for %s" % (len(primary), pchannel))
else:
    message = "No events found for %r in %d seconds of livetime" % (
       pchannel, livetime)
    logger.critical(message)

# cluster primary triggers
clusterkwargs = cp.getparams('primary', 'cluster-')
if clusterkwargs:
    primary = primary.cluster(**clusterkwargs)
    logger.info("%d primary events remain after clustering over %s" %
                (len(primary), clusterkwargs['rank']))

# -- bail out early -----------------------------------------------------------
# the bail out is done here so that we can at least generate the eventual
# configuration file, mainly for HTML purposes

# no segments
if livetime == 0:
    message = ("No active segments found for analysis flag %r in interval "
               "[%d, %d)" % (aflag, start, end))
    logger.critical(message)
    htmlv['context'] = 'info'
    index = html.write_null_page(ifo, start, end, message, **htmlv)
    logger.info("HTML report written to %s" % index)
    sys.exit(0)

# no primary triggers
if len(primary) == 0:
    htmlv['context'] = 'danger'
    index = html.write_null_page(ifo, start, end, message, **htmlv)
    logger.info("HTML report written to %s" % index)
    sys.exit(0)

# otherwise write all primary triggers to ASCII
trigfile = os.path.join(
    trigdir, '%s-HVETO_RAW_TRIGS_ROUND_0-%d-%d.txt' % (ifo, start, duration))
primary.write(trigfile, format='ascii', overwrite=True)

# -- load auxiliary triggers --------------------------------------------------

logger.info("Reading triggers for aux channels...")
counter = multiprocessing.Value('i', 0)

areadkw = cp.getparams('auxiliary', 'read-')
#if acache is not None:  # auto-detect the file format # original line, but not None, we don't have auto-detection
if acache is None:  # auto-detect the file format
    logger.debug('Unsetting the auxiliary trigger file format')
    areadkw['format'] = None
    areadkw['path'] = 'triggers' 
atrigfindkw = cp.getparams('auxiliary', 'trigfind-')


def _get_aux_triggers(channel):
    if acache is None:
        auxcache = None
    else:
        ifo, name = channel.split(':')
        match = "{}-{}".format(ifo, name.replace('-', '_'))
        auxcache = [e for e in acache if Path(e).name.startswith(match)]
    # get triggers
    try:
        trigs = get_triggers(channel, auxetg, analysis.active, snr=minsnr,
                             frange=auxfreq, cache=auxcache, nproc=1,
                             trigfind_kwargs=atrigfindkw, **areadkw)
    # catch error and continue
    except ValueError as e:
        warnings.warn('%s: %s' % (type(e).__name__, str(e)))
        out = None
    else:
        out = (channel, trigs)
    # log result of load
    with counter.get_lock():
        counter.value += 1
        tag = '[%d/%d]' % (counter.value, naux)
        if out is None:  # something went wrong
            logger.critical("    %s Failed to read events for %s"
                            % (tag, channel))
        elif len(trigs):  # no triggers
            logger.debug("    %s Read %d events for %s"
                         % (tag, len(trigs), channel))
        else:  # everything is fine
            logger.warning("    %s No events found for %s"
                           % (tag, channel))
    return out


# map with multiprocessing
if args.nproc > 1:
    pool = multiprocessing.Pool(processes=args.nproc)
    results = pool.map(_get_aux_triggers, auxchannels)
    pool.close()
# map without multiprocessing
else:
    results = map(_get_aux_triggers, auxchannels)

logger.info("All aux events loaded")

auxiliary = dict(x for x in results if x is not None)
auxchannels = sorted(auxiliary.keys())
chanfile = '%s-HVETO_CHANNEL_LIST-%d-%d.txt' % (ifo, start, duration)
with open(chanfile, 'w') as f:
    for chan in auxchannels:
        print(chan, file=f)
logger.info("Recorded list of valid auxiliary channels in %s" % chanfile)

# -- execute hveto analysis ---------------------------------------------------

minsig = cp.getfloat('hveto', 'minimum-significance')
significance = 1e9

pevents = [primary]
pvetoed = []

auxfcol, auxscol = auxiliary[auxchannels[0]].dtype.names[1:3]
slabel = get_column_label(scol)
flabel = get_column_label(fcol)
auxslabel = get_column_label(auxscol)
auxflabel = get_column_label(auxfcol)

rounds = []
round = core.HvetoRound(1, pchannel, rank=scol)
round.segments = analysis.active

while True:
    logger.info("-- Processing round %d --" % round.n)

    # write segments for this round
    segfile = os.path.join(
        segdir, '%s-HVETO_ANALYSIS_SEGS_ROUND_%d-%d-%d.txt'
                % (ifo, round.n, start, duration))
    write_ascii_segments(segfile, round.segments)

    # calculate significances for this round
    if args.nproc > 1:  # multiprocessing
        def _find_max_significance(channels):
            aux = dict((c, auxiliary[c]) for c in channels)
            return core.find_max_significance(primary, aux, pchannel,
                                              snrs, windows, round.livetime)
        # separate channel list into chunks and process each chunk
        pool = multiprocessing.Pool(
            processes=min(args.nproc, len(auxiliary.keys())))
        chunks = utils.channel_groups(list(auxiliary.keys()), args.nproc)
        results = pool.map(_find_max_significance, chunks)
        pool.close()
        winners, sigsets = zip(*results)
        # find winner of chunk winners
        winner = sorted(winners, key=lambda w: w.significance)[-1]
        # flatten sets of significances into one list
        newsignificances = sigsets[0]
        for subdict in sigsets[1:]:
            newsignificances.update(subdict)
    else:  # single process
        winner, newsignificances = core.find_max_significance(
            primary, auxiliary, pchannel, snrs, windows, round.livetime)

    logger.info("Round %d winner: %s" % (round.n, winner.name))

    # plot significance drop here for the last round
    #   only now do we actually have the new data to calculate significance
    #   drop
    if round.n > 1:
        svg = (pngname % 'SIG_DROP').replace('.png', '.svg')  # noqa: F821
        plot.significance_drop(
            svg, oldsignificances, newsignificances,  # noqa: F821
            title=' | '.join([title, subtitle]),  # noqa: F821
            bbox_inches='tight')
        logger.debug("Figure written to %s" % svg)
        svg = FancyPlot(svg, caption=ROUND_CAPTION['SIG_DROP'])
        rounds[-1].plots.append(svg)
    oldsignificances = newsignificances

    # break out of the loop if the significance is below the stopping point
    if winner.significance < minsig:
        logger.info("Maximum signifiance below stopping point")
        logger.debug("    (%.2f < %.2f)" % (winner.significance, minsig))
        logger.info("-- Rounds complete! --")
        break

    # work out the vetoes for this round
    allaux = auxiliary[winner.name][
        auxiliary[winner.name][auxscol] >= winner.snr]
    winner.events = allaux
    coincs = allaux[core.find_coincidences(allaux['time'], primary['time'],
                                           dt=winner.window)]
    round.vetoes = winner.get_segments(allaux['time'])
    flag = DataQualityFlag(
        '%s:HVT-ROUND_%d:1' % (ifo, round.n), active=round.vetoes,
        known=round.segments,
        description="winner=%s, window=%s, snr=%s" % (
            winner.name, winner.window, winner.snr))
    segments[flag.name] = flag
    logger.debug("Generated veto segments for round %d" % round.n)

    # link events before veto for plotting
    before = primary
    beforeaux = auxiliary[winner.name]

    # apply vetoes to primary
    primary, vetoed = core.veto(primary, round.vetoes)
    pevents.append(primary)
    pvetoed.append(vetoed)
    logger.debug("Applied vetoes to primary")

    # record results
    round.winner = winner
    round.efficiency = (len(vetoed), len(primary) + len(vetoed))
    round.use_percentage = (len(coincs), len(winner.events))
    if round.n > 1:
        round.cum_efficiency = (
            len(vetoed) + rounds[-1].cum_efficiency[0],
            rounds[0].efficiency[1])
        round.cum_deadtime = (
            round.deadtime[0] + rounds[-1].cum_deadtime[0],
            livetime)
    else:
        round.cum_efficiency = round.efficiency
        round.cum_deadtime = round.deadtime

    # apply vetoes to auxiliary
    if args.nproc > 1:  # multiprocess
        def _veto(channels):
            return core.veto_all(dict((c, auxiliary[c]) for c in channels),
                                 round.vetoes)
        # separate channel list into chunks and process each chunk
        pool = multiprocessing.Pool(
            processes=min(args.nproc, len(auxiliary.keys())))
        chunks = utils.channel_groups(list(auxiliary.keys()), args.nproc)
        results = pool.map(_veto, chunks)
        pool.close()
        auxiliary = results[0]
        for subdict in results[1:]:
            auxiliary.update(subdict)
    else:  # single process
        auxiliary = core.veto_all(auxiliary, round.vetoes)
    logger.debug("Applied vetoes to auxiliary channels")

    # log results
    logger.info("""Results for round %d:\n\n
winner :          %s
significance :    %s
mu :              %s
snr :             %s
dt :              %s
use_percentage :  %s
efficiency :      %s
deadtime :        %s
cum. efficiency : %s
cum. deadtime :   %s\n\n""" % (
        round.n, round.winner.name, round.winner.significance, round.winner.mu,
        round.winner.snr, round.winner.window, round.use_percentage,
        round.efficiency, round.deadtime, round.cum_efficiency,
        round.cum_deadtime))

    # write segments
    segfile = os.path.join(segdir, '%s-HVETO_VETO_SEGS_ROUND_%d-%d-%d.txt' % (
        ifo, round.n, start, duration))
    write_ascii_segments(segfile, round.vetoes)
    logger.debug("Round %d vetoes written to %s" % (round.n, segfile))
    round.files['VETO_SEGS'] = (segfile,)
    # write triggers
    trigfile = os.path.join(trigdir, '%s-HVETO_%%s_TRIGS_ROUND_%d-%d-%d.txt'
                                     % (ifo, round.n, start, duration))
    for tag, arr in zip(
            ['WINNER', 'VETOED', 'RAW'],
            [winner.events, vetoed, primary]):
        f = trigfile % tag
        arr.write(f, format='ascii', overwrite=True)
        logger.debug("Round %d %s events written to %s"
                     % (round.n, tag.lower(), f))
        round.files[tag] = f

    # record times to omega scan
    if args.omega_scans:
        N = len(vetoed)
        ind = random.sample(range(0, N), min(args.omega_scans, N))
        round.scans = vetoed[ind]
        logger.debug("Collected %d events to omega scan:\n\n%s\n\n"
                     % (len(round.scans), round.scans))

    # -- make some plots --

    pngname = os.path.join(plotdir, '%s-HVETO_%%s_ROUND_%d-%d-%d.png' % (
        ifo, round.n, start, duration))
    wname = texify(round.winner.name)
    beforel = 'Before\n[%d]' % len(before)
    afterl = 'After\n[%d]' % len(primary)
    vetoedl = 'Vetoed\n(primary)\n[%d]' % len(vetoed)
    beforeauxl = 'All\n[%d]' % len(beforeaux)
    usedl = 'Used\n(aux)\n[%d]' % len(winner.events)
    coincl = 'Coinc.\n[%d]' % len(coincs)
    title = '%s Hveto round %d' % (ifo, round.n)
    ptitle = '%s: primary impact' % title
    atitle = '%s: auxiliary use' % title
    subtitle = 'winner: %s [%d-%d]' % (wname, start, end)

    # before/after histogram
    png = pngname % 'HISTOGRAM'
    plot.before_after_histogram(
        png, before[scol], primary[scol],
        label1=beforel, label2=afterl, xlabel=slabel,
        title=ptitle, subtitle=subtitle)
    logger.debug("Figure written to %s" % png)
    png = FancyPlot(png, caption=ROUND_CAPTION['HISTOGRAM'])
    round.plots.append(png)

    # snr versus time
    png = pngname % 'SNR_TIME'
    plot.veto_scatter(
        png, before, vetoed, x='time', y=scol, label1=beforel, label2=vetoedl,
        epoch=start, xlim=[start, end], ylabel=slabel,
        title=ptitle, subtitle=subtitle, legend_title="Primary:")
    logger.debug("Figure written to %s" % png)
    png = FancyPlot(png, caption=ROUND_CAPTION['SNR_TIME'])
    round.plots.append(png)

    # snr versus frequency
    png = pngname % 'SNR_%s' % fcol.upper()
    plot.veto_scatter(
        png, before, vetoed, x=fcol, y=scol, label1=beforel,
        label2=vetoedl, xlabel=flabel, ylabel=slabel, xlim=pfreq,
        title=ptitle, subtitle=subtitle, legend_title="Primary:")
    logger.debug("Figure written to %s" % png)
    png = FancyPlot(png, caption=ROUND_CAPTION['SNR'])
    round.plots.append(png)

    # frequency versus time coloured by SNR
    png = pngname % '%s_TIME' % fcol.upper()
    plot.veto_scatter(
        png, before, vetoed, x='time', y=fcol, color=scol,
        label1=None, label2=None, ylabel=flabel,
        clabel=slabel, clim=[3, 100], cmap='YlGnBu',
        epoch=start, xlim=[start, end], ylim=pfreq,
        title=ptitle, subtitle=subtitle)
    logger.debug("Figure written to %s" % png)
    png = FancyPlot(png, caption=ROUND_CAPTION['TIME'])
    round.plots.append(png)

    # aux used versus frequency
    png = pngname % 'USED_SNR_TIME'
    plot.veto_scatter(
        png, winner.events, vetoed, x='time', y=[auxscol, scol], label1=usedl,
        label2=vetoedl, ylabel=slabel, epoch=start, xlim=[start, end],
        title=atitle, subtitle=subtitle)
    logger.debug("Figure written to %s" % png)
    png = FancyPlot(png, caption=ROUND_CAPTION['USED_SNR_TIME'])
    round.plots.append(png)

    # snr versus time
    png = pngname % 'AUX_SNR_TIME'
    plot.veto_scatter(
        png, beforeaux, (winner.events, coincs), x='time', y=auxscol,
        label1=beforeauxl, label2=(usedl, coincl), epoch=start,
        xlim=[start, end], ylabel=auxslabel, title=atitle, subtitle=subtitle)
    logger.debug("Figure written to %s" % png)
    png = FancyPlot(png, caption=ROUND_CAPTION['AUX_SNR_TIME'])
    round.plots.append(png)

    # snr versus frequency
    png = pngname % 'AUX_SNR_FREQUENCY'
    plot.veto_scatter(
        png, beforeaux, (winner.events, coincs), x=auxfcol, y=auxscol,
        label1=beforeauxl, label2=(usedl, coincl), xlabel=auxflabel,
        ylabel=auxslabel, title=atitle, subtitle=subtitle, legend_title="Aux:")
    logger.debug("Figure written to %s" % png)
    png = FancyPlot(png, caption=ROUND_CAPTION['AUX_SNR_FREQUENCY'])
    round.plots.append(png)

    # frequency versus time coloured by SNR
    png = pngname % 'AUX_FREQUENCY_TIME'
    plot.veto_scatter(
        png, beforeaux, (winner.events, coincs), x='time', y=auxfcol,
        color=auxscol, label1=None, label2=[None, None], ylabel=auxflabel,
        clabel=auxslabel, clim=[3, 100], cmap='YlGnBu', epoch=start,
        xlim=[start, end], title=atitle, subtitle=subtitle)
    logger.debug("Figure written to %s" % png)
    png = FancyPlot(png, caption=ROUND_CAPTION['AUX_FREQUENCY_TIME'])
    round.plots.append(png)

    # move to the next round
    rounds.append(round)
    round = core.HvetoRound(round.n + 1, pchannel, rank=scol,
                            segments=round.segments-round.vetoes)

# write file with all segments
segfile = os.path.join(
    segdir, '%s-HVETO_SEGMENTS-%d-%d.h5' % (ifo, start, duration))
segments.write(segfile, overwrite=True)
logger.debug("Segment summary written to %s" % segfile)

logger.debug("Making summary figures...")

# -- exit early if no rounds above threshold

if not rounds:
    message = ("No rounds completed above threshold. Analysis stopped with "
               "%s achieving significance of %.2f"
               % (winner.name, winner.significance))
    logger.critical(message)
    message = message.replace(
        winner.name, cis_link(winner.name, class_='alert-link'))
    message += '<br>[T<sub>win</sub>: %ss, SNR: %s]' % (
        winner.window, winner.snr)
    htmlv['context'] = 'warning'
    index = html.write_null_page(ifo, start, end, message, **htmlv)
    logger.info("HTML report written to %s" % index)
    sys.exit(0)

# -- plot all rounds impact
pngname = os.path.join(plotdir, '%s-HVETO_%%s_ALL_ROUNDS-%d-%d.png' % (
    ifo, start, duration))
plots = []
title = '%s Hveto all rounds' % args.ifo
subtitle = '%d rounds | %d-%d' % (len(rounds), start, end)

# before/after histogram
png = pngname % 'HISTOGRAM'
beforel = 'Before analysis [%d events]' % len(pevents[0])
afterl = 'After %d rounds [%d]' % (len(pevents) - 1, len(pevents[-1]))
plot.before_after_histogram(
    png, pevents[0][scol], pevents[-1][scol],
    label1=beforel, label2=afterl, xlabel=slabel,
    title=title, subtitle=subtitle)
png = FancyPlot(png, caption=HEADER_CAPTION['HISTOGRAM'])
plots.append(png)
logger.debug("Figure written to %s" % png)

# efficiency/deadtime curve
png = pngname % 'ROC'
plot.hveto_roc(png, rounds, title=title, subtitle=subtitle)
png = FancyPlot(png, caption=HEADER_CAPTION['ROC'])
plots.append(png)
logger.debug("Figure written to %s" % png)

# frequency versus time
png = pngname % '%s_TIME' % fcol.upper()
labels = [str(r.n) for r in rounds]
legtitle = 'Vetoed at\nround'
plot.veto_scatter(
    png, pevents[0], pvetoed,
    label1='', label2=labels, title=title,
    subtitle=subtitle, ylabel=flabel, x='time', y=fcol,
    epoch=start, xlim=[start, end], legend_title=legtitle)
png = FancyPlot(png, caption=HEADER_CAPTION['TIME'])
plots.append(png)
logger.debug("Figure written to %s" % png)

# snr versus time
png = pngname % 'SNR_TIME'
plot.veto_scatter(
    png, pevents[0], pvetoed, label1='', label2=labels, title=title,
    subtitle=subtitle, ylabel=slabel, x='time', y=scol,
    epoch=start, xlim=[start, end], legend_title=legtitle)
png = FancyPlot(png, caption=HEADER_CAPTION['SNR_TIME'])
plots.append(png)
logger.debug("Figure written to %s" % png)

# -- write summary states to ASCII table and JSON
json_ = {
    'user': getuser(),
    'host': getfqdn(),
    'date': str(datetime.datetime.now()),
    'configuration': inifile,
    'ifo': ifo,
    'gpsstart': start,
    'gpsend': end,
    'call': ' '.join(sys.argv),
    'rounds': [],
}
with open('summary-stats.txt', 'w') as f:
    # print header
    print('#N winner window SNR significance nveto use-percentage efficiency '
          'deadtime cumulative-efficiency cumulative-deadtime', file=f)
    for r in rounds:
        # extract relevant statistics
        results = [
            ('round', r.n),
            ('name', r.winner.name),
            ('window', r.winner.window),
            ('snr', r.winner.snr),
            ('significance', r.winner.significance),
            ('nveto', r.efficiency[0]),
            ('use-percentage',
                r.use_percentage[0] / r.use_percentage[1] * 100.),
            ('efficiency', r.efficiency[0] / r.efficiency[1] * 100.),
            ('deadtime', r.deadtime[0] / r.deadtime[1] * 100.),
            ('cumulative-efficiency',
                r.cum_efficiency[0] / r.cum_efficiency[1] * 100.),
            ('cumulative-deadtime',
                r.cum_deadtime[0] / r.cum_deadtime[1] * 100.),
        ]
        # write to ASCII
        print(' '.join(map(str, list(zip(*results))[1])), file=f)
        # write to JSON
        results.append(('files', r.files))
        json_['rounds'].append(dict(results))
logger.debug("Summary table written to %s" % f.name)

with open('summary-stats.json', 'w') as f:
    json.dump(json_, f, sort_keys=True)
logger.debug("Summary JSON written to %s" % f.name)

# -- generate workflow for omega scans

if args.omega_scans:
    omegatimes = list(map(str, sorted(numpy.unique(
        [t['time'] for r in rounds for t in r.scans]))))
    logger.debug("Collected %d times to omega scan" % len(omegatimes))
    newtimes = [t for t in omegatimes if not
                os.path.exists(os.path.join(omegadir, str(t)))]
    logger.debug("%d scans already complete or in progress, %d remaining"
                 % (len(omegatimes) - len(newtimes), len(newtimes)))
    if len(newtimes) > 0:
        logger.info('Creating workflow for omega scans')
        flags = batch.get_command_line_flags(
            ifo=ifo,
            ignore_state_flags=True)
        condorcmds = batch.get_condor_arguments(
            timeout=4,
            gps=start)
        dagman = batch.generate_dag(
            newtimes,
            flags=flags,
            submit=True,
            outdir=omegadir,
            condor_commands=condorcmds)
        logger.info('Launched {} omega scans to condor'.format(len(newtimes)))
    else:
        logger.debug('Skipping omega scans')

# -- write HTML and finish

index = html.write_hveto_page(
    ifo, start, end, rounds, plots,
    winners=[r.winner.name for r in rounds], **htmlv)
logger.debug("HTML written to %s" % index)
logger.debug("Analysis completed in %d seconds" % (time.time() - jobstart))
logger.info("-- Hveto complete --")