#!/usr/bin/python3 -s
# coding=utf-8
# Copyright (C) LIGO Scientific Collaboration (2015-)
#
# This file is part of the GW DetChar python package.
#
# GW DetChar 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.
#
# GW DetChar 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 GW DetChar.  If not, see <http://www.gnu.org/licenses/>.

"""Batch-generate a series of Omega-pipeline scans.

GPS times can be given individually on the command-line, one after the other,
or can be bundled into one file formatted where the first column contains
the GPS times (other columns are ignored).

The output of this script is a condor workflow in the form of a DAG file,
with associated condor submit (`.sub`) file in the output directory.
Submitting the workflow to Condor will result in the scans being processed
in parallel.
"""

import os
import subprocess

from gwdetchar import cli
from gwdetchar.omega import batch

__author__ = 'Duncan Macleod <duncan.macleod@ligo.org'
__credits__ = 'Alex Urban <alexander.urban@ligo.org'

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

parser = cli.create_parser(description=__doc__)

parser.add_argument('gps-time', nargs='+',
                    help='GPS time(s) to scan, or path to a file '
                         'containing a single column of such times')
cli.add_ifo_option(parser)
parser.add_argument('-o', '--output-dir', default=os.getcwd(),
                    help='output directory for all scans')

parser.add_argument(
    '-f', '--config-file',
    help='path to configuration file to use, default: '
         'choose based on observatory, epoch, and pipeline')
parser.add_argument('--colormap', default='viridis',
                    help='name of colormap to use, default: %(default)s')
parser.add_argument('-d', '--disable-correlation', action='store_true',
                    default=False, help='disable cross-correlation of aux '
                                        'channels, default: False')
parser.add_argument('-D', '--disable-checkpoint', action='store_true',
                    default=False, help='disable checkpointing from previous '
                                        'runs, default: False')
parser.add_argument('-t', '--far-threshold', type=float, default=3.171e-8,
                    help='white noise false alarm rate threshold (Hz) for '
                         'processing channels, default: %(default)s')
parser.add_argument('-s', '--ignore-state-flags', action='store_true',
                    default=False, help='ignore state flag definitions in '
                                        'the configuration, default: False')
cli.add_nproc_option(parser)

cargs = parser.add_argument_group('Condor options')
cargs.add_argument('-u', '--universe', default='vanilla', type=str,
                   help='universe for condor processing')
cargs.add_argument('--submit', action='store_true', default=False,
                   help='submit DAG directly to condor queue')
cargs.add_argument('--monitor', action='store_true', default=False,
                   help='monitor the DAG progress after submission; '
                        'only used with --submit')
cargs.add_argument('--condor-accounting-group',
                   default=batch.ACCOUNTING_GROUP,
                   help='accounting_group for condor submission on the LIGO '
                        'Data Grid, include \'{epoch}\' (with curly brackets) '
                        'to auto-substitute the appropriate epoch based on '
                        'the GPS times')
cargs.add_argument('--condor-accounting-group-user',
                   default=batch.ACCOUNTING_GROUP_USER,
                   help='accounting_group_user for condor submission on the '
                        'LIGO Data Grid')
cargs.add_argument('--condor-timeout', type=float, default=None, metavar='T',
                   help='configure condor to terminate jobs after T hours '
                        'to prevent idling, default: %(default)s')
cargs.add_argument('--condor-command', action='append', default=[],
                   help="Extra condor submit commands to add to "
                        "gw_summary submit file. Can be given "
                        "multiple times in the form \"key=value\"")

args = parser.parse_args()

# set up output directory
outdir = os.path.abspath(os.path.expanduser(args.output_dir))

# parse times
times = getattr(args, 'gps-time')

if len(times) == 1:
    try:  # try converting to GPS
        times = [float(times[0])]
    except (TypeError, ValueError):  # otherwise read as file
        import numpy
        times = numpy.loadtxt(times[0], dtype=float, ndmin=1)
else:
    times = list(map(float, times))

# get condor arguments
condorcmds = batch.get_condor_arguments(
    accounting_group=args.condor_accounting_group,
    accounting_group_user=args.condor_accounting_group_user,
    timeout=args.condor_timeout,
    extra_commands=args.condor_command,
    gps=max(times),
)

# get command-line flags
flags = batch.get_command_line_flags(
    args.ifo,
    colormap=args.colormap,
    nproc=args.nproc,
    far=args.far_threshold,
    config_file=args.config_file,
    disable_correlation=args.disable_correlation,
    disable_checkpoint=args.disable_checkpoint,
    ignore_state_flags=args.ignore_state_flags,
)

# -- generate workflow --------------------------------------------------------

# write and submit the DAG
dagman = batch.generate_dag(
    times,
    flags=flags,
    tag=os.path.basename(__file__),
    submit=args.submit,
    outdir=outdir,
    universe=args.universe,
    condor_commands=condorcmds,
)

# monitor progress
if args.submit and args.monitor:
    print("Monitoring progress of {0.submit_file}".format(dagman))
    try:
        subprocess.check_call(
            ["pycondor", "monitor", dagman.submit_file],
        )
    except KeyboardInterrupt:
        pass
