#!/usr/bin/env python
# coding: utf-8

# Copyright 2022
# Maximiliano Isi <max.isi@ligo.org>
#
# This program 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 2 of the License, or
# (at your option) any later version.
#
# This program 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 this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA 02110-1301, USA.

import sys
import os
from pylab import *
import ringdown as rd
import argparse
import configparser
import logging
import time
from ast import literal_eval


##############################################################################
# PARSE INPUT
##############################################################################

parser  = argparse.ArgumentParser(description="Analyze event ringdown at "
                                  "different times.")
parser.add_argument('config', help="Path to configuration file.")
parser.add_argument('-o', '--outdir', default=None, help="Output directory.")
parser.add_argument('--force', action='store_true', help="Run even if output "
                    "directory already exists (pre-existing runs will not be "
                    "repeated).")
parser.add_argument('-s', '--submit', action='store_true', help="Submit slurm "
                    "job automatically.")
parser.add_argument('--nodes', default=20, type=int,
                    help="Maximum number of nodes to request through SLURM.")
parser.add_argument('--seed', default=None, type=int, help="Random seed.")
parser.add_argument('-v', '--verbose', action='store_true')

args = parser.parse_args()

# set up logging and procure Git repo hash
if args.verbose:
    logging.getLogger().setLevel(logging.INFO)

try:
    import git
    repo = git.Repo(search_parent_directories=True)
    githash = repo.head.object.hexsha
    logging.info('GIT {}'.format(githash))
except ImportError:
    githash = 'unknown'
    logging.info("install gitpython to report repo hash")
except Exception:
    githash = 'unknown'
    logging.warning("failed to locate git repo")

# load config file 
config_path = os.path.abspath(args.config)
if not os.path.isfile(config_path):
    raise FileNotFoundError(f"unable to load: {config_path}")
logging.info(f"Loading config from {config_path}")
config = configparser.ConfigParser()
config.read(config_path)

# set random seed (purposedly fail if not provided)
seed = args.seed or config.getint('pipe', 'seed')
logging.info("Random seed set to {}".format(seed))
rng = np.random.default_rng(seed)

# determine run directory
outdir = args.outdir or config.get('pipe', 'outdir', fallback=None)
if not outdir:
    outdir = 'ringdown_pipe_{:.0f}'.format(time.time())
    logging.warning("No run dir provided, defaulting to {}".format(outdir))
outdir = os.path.abspath(outdir)
logging.info("Running in {}".format(outdir))

rerun = False
if os.path.exists(outdir):
    if args.force:
        logging.warning("Run directory already exists.")
        rerun = True
    else:
        raise FileExistsError("Run directory already exists. Exiting.")
else:
    os.makedirs(outdir)

PATHS = {
    'config': 'config.ini',
    'command': 'pipe.sh',
    'run_config': 'engine/{t0:.6f}/config.ini',
    'run_result': 'engine/{t0:.6f}/result.nc',
    'run_task': 'TaskFile',
    'exe': 'submit.sh',
    't0': 'start_times.txt',
}
PATHS = {k: os.path.join(outdir, v) for k,v in PATHS.items()}

# record config and arguments for reproducibility
with open(PATHS['config'], 'w') as f:
    config.write(f)

with open(PATHS['command'], 'w') as f:
    f.write("{}\n\n".format(' '.join(sys.argv)))
    for k,v in vars(args).items():
        f.write("# {}: {}\n".format(k,v))
    f.write("\n# {}".format(os.path.realpath(__file__)))
    f.write("\n# GIT hash: {}".format(githash))
    f.write("\n# ringdown v {}".format(getattr(rd, '__version__', 'unknown')))

# Identify target analysis times. There will be three possibilities:
#     1- listing the times explicitly
#     2- listing time differences with respect to a reference time
#     3- providing start, stop, step instructions to construct start times
#        (potentially relative to a reference time)
# Time steps/differences can be specified in seconds or M, if a reference mass
# is provided (in solar masses).

# Define valid options to specify the start times
T0_SECT = "pipe"
T0_KEYS = {
    'ref': 't0-ref',
    'list': 't0-list',
    'delta': 't0-delta-list',
    'step': 't0-step',
    'start': 't0-start',
    'stop': 't0-stop'
}
start_stop_step = [T0_KEYS[k] for k in ['start', 'stop', 'step']]

# First make sure that only compatible t0 options were provided
incompatible_sets = [['ref', 'list'], ['delta', 'list']]
incompatible_sets += [[k, 'delta'] for k in ['start', 'stop', 'step']]
for bad_set in incompatible_sets:
    opt_names = [T0_KEYS[k] for k in bad_set]
    if all([k in config[T0_SECT] for k in opt_names]):
        raise ValueError("incompatible T0 options: {}".format(opt_names))

# Look for a reference mass, to be used when stepping in time
m_ref = config.getfloat(T0_SECT, 'M-ref', fallback=None)
if m_ref:
    # reference time translating from solar masses
    tm_ref = m_ref * rd.qnms.T_MSUN
    logging.info("Reference mass: {} Msun ({} s)".format(m_ref, tm_ref))
else:
    # no reference mass provided, so will default to seconds
    tm_ref = 1

# Look for reference time to be used to construct start times
t0ref = config.getfloat(T0_SECT, T0_KEYS['ref'], fallback=0)

# Now we can safely interpret the options assuming one of three cases
if T0_KEYS['list'] in config[T0_SECT]:
    t0s = array(literal_eval(config.get(T0_SECT, T0_KEYS['list'])))
elif T0_KEYS['delta'] in config[T0_SECT]:
    dt0s = array(literal_eval(config.get(T0_SECT, T0_KEYS['delta'])))
    t0s = dt0s*tm_ref + t0ref
elif any([k in config[T0_SECT] for k in start_stop_step]):
    if not all([k in config[T0_SECT] for k in start_stop_step]):
        missing = [k for k in start_stop_step if k not in config[T0_SECT]]
        raise ValueError("missing start/stop/step options: {}".format(missing))
    # add a safety check here, in case the user mistakenly requests stepping
    # based on a GPS time and provides a reference GPS time
    start, stop, step  = [config.getfloat(T0_SECT, k) for k in start_stop_step]
    if start > 500 and t0ref > 1E8:
        logging.warning("high reference time and stepping start---did you "
                        "accidentally provide GPS times twice?")
    t0s = arange(start, stop, step)*tm_ref + t0ref
else:
    raise ValueError("no timing instructions in [{}] section; valid options "
                     "are: {}".format(T0_SECT, list(T0_KEYS.values())))

# Record start times
savetxt(PATHS["t0"], t0s)

##############################################################################
# SET UP RUNS
##############################################################################

if rerun and os.path.exists(PATHS['run_task']):
    os.remove(PATHS['run_task'])

TASK = "mkdir -p {rundir}; cd {rundir}; ringdown_fit -o {result} {config} "\
       "&> run.log\n"

logging.info(f"Processing {len(t0s)} start times (recorded in {PATHS['t0']})")
for i, t0 in enumerate(t0s):
    # Set up child configuration file
    cpath = PATHS['run_config'].format(i=i, t0=t0)
    rundir = os.path.dirname(cpath)
    os.makedirs(rundir, exist_ok=True)
    config_child = configparser.ConfigParser()
    config_child.read(PATHS['config'])

    # Set start time
    # Restore conditioning section
    config_child['target']['t0'] = str(t0)

    # Write out config file for this run
    rpath = PATHS['run_result'].format(i=i, t0=t0)
    if os.path.exists(rpath):
        logging.info(f"Run {i} already exists. Will skip.")
    else:
        with open(cpath, 'w') as f:
            config_child.write(f)
        with open(PATHS['run_task'], 'a') as f:
            f.write(TASK.format(rundir=rundir, result=rpath, config=cpath))

logging.info("Done processing {} runs.".format(len(t0s)))


##############################################################################
# SLURM WORKFLOW
##############################################################################

# rusty Broadwell nodes have 28 cores, so each can fit 7 4-core jobs at a time
# given that, make the number of jobs executed at a time a multiple of 7
NSLURM = 7*args.nodes
NCPU = 4

# a user may allocate up to 1280 cores at one time in a center partition, so
# issue a warning if this limit will be exceeded
if NSLURM*NCPU > 1280:
    w = "Requested number of cores ({}) above 1280 user limit."
    logging.warning(w.format(NSLURM*NCPU))

EXE = [
    '#! /usr/bin/env bash',
    'cd {}'.format(outdir),
    'sbatch -p cca -n {} -c {} disBatch {}'.format(NSLURM, NCPU, PATHS['run_task']),
    'cd -',
]

epath = PATHS['exe']
with open(epath, 'w') as f:
    f.write('\n'.join(EXE))
st = os.stat(epath)
os.chmod(epath, st.st_mode | 0o111)
if args.submit:
    print("Submitting: {}".format(epath))
    import subprocess
    subprocess.run(epath)
else:
    print("Submit by running: {}".format(epath))

625f3f56e3598e2951dafc52a4a67de856d78611
<mask token>

<mask token>
parser.add_argument('config', help='Path to configuration file.')
parser.add_argument('-o', '--outdir', default=None, help='Output directory.')
parser.add_argument('--force', action='store_true', help=
    'Run even if output directory already exists (pre-existing runs will not be repeated).'
    )
parser.add_argument('-s', '--submit', action='store_true', help=
    'Submit slurm job automatically.')
parser.add_argument('--nodes', default=20, type=int, help=
    'Maximum number of nodes to request through SLURM.')
parser.add_argument('--seed', default=None, type=int, help='Random seed.')
parser.add_argument('-v', '--verbose', action='store_true')
<mask token>
if args.verbose:
    logging.getLogger().setLevel(logging.INFO)
try:
    import git
    repo = git.Repo(search_parent_directories=True)
    githash = repo.head.object.hexsha
    logging.info('GIT {}'.format(githash))
except ImportError:
    githash = 'unknown'
    logging.info('install gitpython to report repo hash')
except Exception:
    githash = 'unknown'
    logging.warning('failed to locate git repo')
<mask token>
if not os.path.isfile(config_path):
    raise FileNotFoundError(f'unable to load: {config_path}')
logging.info(f'Loading config from {config_path}')
<mask token>
config.read(config_path)
<mask token>
logging.info('Random seed set to {}'.format(seed))
<mask token>
if not outdir:
    outdir = 'ringdown_pipe_{:.0f}'.format(time.time())
    logging.warning('No run dir provided, defaulting to {}'.format(outdir))
<mask token>
logging.info('Running in {}'.format(outdir))
<mask token>
if os.path.exists(outdir):
    if args.force:
        logging.warning('Run directory already exists.')
        rerun = True
    else:
        raise FileExistsError('Run directory already exists. Exiting.')
else:
    os.makedirs(outdir)
<mask token>
with open(PATHS['config'], 'w') as f:
    config.write(f)
with open(PATHS['command'], 'w') as f:
    f.write('{}\n\n'.format(' '.join(sys.argv)))
    for k, v in vars(args).items():
        f.write('# {}: {}\n'.format(k, v))
    f.write('\n# {}'.format(os.path.realpath(__file__)))
    f.write('\n# GIT hash: {}'.format(githash))
    f.write('\n# ringdown v {}'.format(getattr(rd, '__version__', 'unknown')))
<mask token>
incompatible_sets += [[k, 'delta'] for k in ['start', 'stop', 'step']]
for bad_set in incompatible_sets:
    opt_names = [T0_KEYS[k] for k in bad_set]
    if all([(k in config[T0_SECT]) for k in opt_names]):
        raise ValueError('incompatible T0 options: {}'.format(opt_names))
<mask token>
if m_ref:
    tm_ref = m_ref * rd.qnms.T_MSUN
    logging.info('Reference mass: {} Msun ({} s)'.format(m_ref, tm_ref))
else:
    tm_ref = 1
<mask token>
if T0_KEYS['list'] in config[T0_SECT]:
    t0s = array(literal_eval(config.get(T0_SECT, T0_KEYS['list'])))
elif T0_KEYS['delta'] in config[T0_SECT]:
    dt0s = array(literal_eval(config.get(T0_SECT, T0_KEYS['delta'])))
    t0s = dt0s * tm_ref + t0ref
elif any([(k in config[T0_SECT]) for k in start_stop_step]):
    if not all([(k in config[T0_SECT]) for k in start_stop_step]):
        missing = [k for k in start_stop_step if k not in config[T0_SECT]]
        raise ValueError('missing start/stop/step options: {}'.format(missing))
    start, stop, step = [config.getfloat(T0_SECT, k) for k in start_stop_step]
    if start > 500 and t0ref > 100000000.0:
        logging.warning(
            'high reference time and stepping start---did you accidentally provide GPS times twice?'
            )
    t0s = arange(start, stop, step) * tm_ref + t0ref
else:
    raise ValueError(
        'no timing instructions in [{}] section; valid options are: {}'.
        format(T0_SECT, list(T0_KEYS.values())))
savetxt(PATHS['t0'], t0s)
if rerun and os.path.exists(PATHS['run_task']):
    os.remove(PATHS['run_task'])
<mask token>
logging.info(f"Processing {len(t0s)} start times (recorded in {PATHS['t0']})")
for i, t0 in enumerate(t0s):
    cpath = PATHS['run_config'].format(i=i, t0=t0)
    rundir = os.path.dirname(cpath)
    os.makedirs(rundir, exist_ok=True)
    config_child = configparser.ConfigParser()
    config_child.read(PATHS['config'])
    config_child['target']['t0'] = str(t0)
    rpath = PATHS['run_result'].format(i=i, t0=t0)
    if os.path.exists(rpath):
        logging.info(f'Run {i} already exists. Will skip.')
    else:
        with open(cpath, 'w') as f:
            config_child.write(f)
        with open(PATHS['run_task'], 'a') as f:
            f.write(TASK.format(rundir=rundir, result=rpath, config=cpath))
logging.info('Done processing {} runs.'.format(len(t0s)))
<mask token>
if NSLURM * NCPU > 1280:
    w = 'Requested number of cores ({}) above 1280 user limit.'
    logging.warning(w.format(NSLURM * NCPU))
<mask token>
with open(epath, 'w') as f:
    f.write('\n'.join(EXE))
<mask token>
os.chmod(epath, st.st_mode | 73)
if args.submit:
    print('Submitting: {}'.format(epath))
    import subprocess
    subprocess.run(epath)
else:
    print('Submit by running: {}'.format(epath))

<mask token>
parser = argparse.ArgumentParser(description=
    'Analyze event ringdown at different times.')
parser.add_argument('config', help='Path to configuration file.')
parser.add_argument('-o', '--outdir', default=None, help='Output directory.')
parser.add_argument('--force', action='store_true', help=
    'Run even if output directory already exists (pre-existing runs will not be repeated).'
    )
parser.add_argument('-s', '--submit', action='store_true', help=
    'Submit slurm job automatically.')
parser.add_argument('--nodes', default=20, type=int, help=
    'Maximum number of nodes to request through SLURM.')
parser.add_argument('--seed', default=None, type=int, help='Random seed.')
parser.add_argument('-v', '--verbose', action='store_true')
args = parser.parse_args()
if args.verbose:
    logging.getLogger().setLevel(logging.INFO)
try:
    import git
    repo = git.Repo(search_parent_directories=True)
    githash = repo.head.object.hexsha
    logging.info('GIT {}'.format(githash))
except ImportError:
    githash = 'unknown'
    logging.info('install gitpython to report repo hash')
except Exception:
    githash = 'unknown'
    logging.warning('failed to locate git repo')
config_path = os.path.abspath(args.config)
if not os.path.isfile(config_path):
    raise FileNotFoundError(f'unable to load: {config_path}')
logging.info(f'Loading config from {config_path}')
config = configparser.ConfigParser()
config.read(config_path)
seed = args.seed or config.getint('pipe', 'seed')
logging.info('Random seed set to {}'.format(seed))
rng = np.random.default_rng(seed)
outdir = args.outdir or config.get('pipe', 'outdir', fallback=None)
if not outdir:
    outdir = 'ringdown_pipe_{:.0f}'.format(time.time())
    logging.warning('No run dir provided, defaulting to {}'.format(outdir))
outdir = os.path.abspath(outdir)
logging.info('Running in {}'.format(outdir))
rerun = False
if os.path.exists(outdir):
    if args.force:
        logging.warning('Run directory already exists.')
        rerun = True
    else:
        raise FileExistsError('Run directory already exists. Exiting.')
else:
    os.makedirs(outdir)
PATHS = {'config': 'config.ini', 'command': 'pipe.sh', 'run_config':
    'engine/{t0:.6f}/config.ini', 'run_result': 'engine/{t0:.6f}/result.nc',
    'run_task': 'TaskFile', 'exe': 'submit.sh', 't0': 'start_times.txt'}
PATHS = {k: os.path.join(outdir, v) for k, v in PATHS.items()}
with open(PATHS['config'], 'w') as f:
    config.write(f)
with open(PATHS['command'], 'w') as f:
    f.write('{}\n\n'.format(' '.join(sys.argv)))
    for k, v in vars(args).items():
        f.write('# {}: {}\n'.format(k, v))
    f.write('\n# {}'.format(os.path.realpath(__file__)))
    f.write('\n# GIT hash: {}'.format(githash))
    f.write('\n# ringdown v {}'.format(getattr(rd, '__version__', 'unknown')))
T0_SECT = 'pipe'
T0_KEYS = {'ref': 't0-ref', 'list': 't0-list', 'delta': 't0-delta-list',
    'step': 't0-step', 'start': 't0-start', 'stop': 't0-stop'}
start_stop_step = [T0_KEYS[k] for k in ['start', 'stop', 'step']]
incompatible_sets = [['ref', 'list'], ['delta', 'list']]
incompatible_sets += [[k, 'delta'] for k in ['start', 'stop', 'step']]
for bad_set in incompatible_sets:
    opt_names = [T0_KEYS[k] for k in bad_set]
    if all([(k in config[T0_SECT]) for k in opt_names]):
        raise ValueError('incompatible T0 options: {}'.format(opt_names))
m_ref = config.getfloat(T0_SECT, 'M-ref', fallback=None)
if m_ref:
    tm_ref = m_ref * rd.qnms.T_MSUN
    logging.info('Reference mass: {} Msun ({} s)'.format(m_ref, tm_ref))
else:
    tm_ref = 1
t0ref = config.getfloat(T0_SECT, T0_KEYS['ref'], fallback=0)
if T0_KEYS['list'] in config[T0_SECT]:
    t0s = array(literal_eval(config.get(T0_SECT, T0_KEYS['list'])))
elif T0_KEYS['delta'] in config[T0_SECT]:
    dt0s = array(literal_eval(config.get(T0_SECT, T0_KEYS['delta'])))
    t0s = dt0s * tm_ref + t0ref
elif any([(k in config[T0_SECT]) for k in start_stop_step]):
    if not all([(k in config[T0_SECT]) for k in start_stop_step]):
        missing = [k for k in start_stop_step if k not in config[T0_SECT]]
        raise ValueError('missing start/stop/step options: {}'.format(missing))
    start, stop, step = [config.getfloat(T0_SECT, k) for k in start_stop_step]
    if start > 500 and t0ref > 100000000.0:
        logging.warning(
            'high reference time and stepping start---did you accidentally provide GPS times twice?'
            )
    t0s = arange(start, stop, step) * tm_ref + t0ref
else:
    raise ValueError(
        'no timing instructions in [{}] section; valid options are: {}'.
        format(T0_SECT, list(T0_KEYS.values())))
savetxt(PATHS['t0'], t0s)
if rerun and os.path.exists(PATHS['run_task']):
    os.remove(PATHS['run_task'])
TASK = (
    'mkdir -p {rundir}; cd {rundir}; ringdown_fit -o {result} {config} &> run.log\n'
    )
logging.info(f"Processing {len(t0s)} start times (recorded in {PATHS['t0']})")
for i, t0 in enumerate(t0s):
    cpath = PATHS['run_config'].format(i=i, t0=t0)
    rundir = os.path.dirname(cpath)
    os.makedirs(rundir, exist_ok=True)
    config_child = configparser.ConfigParser()
    config_child.read(PATHS['config'])
    config_child['target']['t0'] = str(t0)
    rpath = PATHS['run_result'].format(i=i, t0=t0)
    if os.path.exists(rpath):
        logging.info(f'Run {i} already exists. Will skip.')
    else:
        with open(cpath, 'w') as f:
            config_child.write(f)
        with open(PATHS['run_task'], 'a') as f:
            f.write(TASK.format(rundir=rundir, result=rpath, config=cpath))
logging.info('Done processing {} runs.'.format(len(t0s)))
NSLURM = 7 * args.nodes
NCPU = 4
if NSLURM * NCPU > 1280:
    w = 'Requested number of cores ({}) above 1280 user limit.'
    logging.warning(w.format(NSLURM * NCPU))
EXE = ['#! /usr/bin/env bash', 'cd {}'.format(outdir),
    'sbatch -p cca -n {} -c {} disBatch {}'.format(NSLURM, NCPU, PATHS[
    'run_task']), 'cd -']
epath = PATHS['exe']
with open(epath, 'w') as f:
    f.write('\n'.join(EXE))
st = os.stat(epath)
os.chmod(epath, st.st_mode | 73)
if args.submit:
    print('Submitting: {}'.format(epath))
    import subprocess
    subprocess.run(epath)
else:
    print('Submit by running: {}'.format(epath))

import sys
import os
from pylab import *
import ringdown as rd
import argparse
import configparser
import logging
import time
from ast import literal_eval
parser = argparse.ArgumentParser(description=
    'Analyze event ringdown at different times.')
parser.add_argument('config', help='Path to configuration file.')
parser.add_argument('-o', '--outdir', default=None, help='Output directory.')
parser.add_argument('--force', action='store_true', help=
    'Run even if output directory already exists (pre-existing runs will not be repeated).'
    )
parser.add_argument('-s', '--submit', action='store_true', help=
    'Submit slurm job automatically.')
parser.add_argument('--nodes', default=20, type=int, help=
    'Maximum number of nodes to request through SLURM.')
parser.add_argument('--seed', default=None, type=int, help='Random seed.')
parser.add_argument('-v', '--verbose', action='store_true')
args = parser.parse_args()
if args.verbose:
    logging.getLogger().setLevel(logging.INFO)
try:
    import git
    repo = git.Repo(search_parent_directories=True)
    githash = repo.head.object.hexsha
    logging.info('GIT {}'.format(githash))
except ImportError:
    githash = 'unknown'
    logging.info('install gitpython to report repo hash')
except Exception:
    githash = 'unknown'
    logging.warning('failed to locate git repo')
config_path = os.path.abspath(args.config)
if not os.path.isfile(config_path):
    raise FileNotFoundError(f'unable to load: {config_path}')
logging.info(f'Loading config from {config_path}')
config = configparser.ConfigParser()
config.read(config_path)
seed = args.seed or config.getint('pipe', 'seed')
logging.info('Random seed set to {}'.format(seed))
rng = np.random.default_rng(seed)
outdir = args.outdir or config.get('pipe', 'outdir', fallback=None)
if not outdir:
    outdir = 'ringdown_pipe_{:.0f}'.format(time.time())
    logging.warning('No run dir provided, defaulting to {}'.format(outdir))
outdir = os.path.abspath(outdir)
logging.info('Running in {}'.format(outdir))
rerun = False
if os.path.exists(outdir):
    if args.force:
        logging.warning('Run directory already exists.')
        rerun = True
    else:
        raise FileExistsError('Run directory already exists. Exiting.')
else:
    os.makedirs(outdir)
PATHS = {'config': 'config.ini', 'command': 'pipe.sh', 'run_config':
    'engine/{t0:.6f}/config.ini', 'run_result': 'engine/{t0:.6f}/result.nc',
    'run_task': 'TaskFile', 'exe': 'submit.sh', 't0': 'start_times.txt'}
PATHS = {k: os.path.join(outdir, v) for k, v in PATHS.items()}
with open(PATHS['config'], 'w') as f:
    config.write(f)
with open(PATHS['command'], 'w') as f:
    f.write('{}\n\n'.format(' '.join(sys.argv)))
    for k, v in vars(args).items():
        f.write('# {}: {}\n'.format(k, v))
    f.write('\n# {}'.format(os.path.realpath(__file__)))
    f.write('\n# GIT hash: {}'.format(githash))
    f.write('\n# ringdown v {}'.format(getattr(rd, '__version__', 'unknown')))
T0_SECT = 'pipe'
T0_KEYS = {'ref': 't0-ref', 'list': 't0-list', 'delta': 't0-delta-list',
    'step': 't0-step', 'start': 't0-start', 'stop': 't0-stop'}
start_stop_step = [T0_KEYS[k] for k in ['start', 'stop', 'step']]
incompatible_sets = [['ref', 'list'], ['delta', 'list']]
incompatible_sets += [[k, 'delta'] for k in ['start', 'stop', 'step']]
for bad_set in incompatible_sets:
    opt_names = [T0_KEYS[k] for k in bad_set]
    if all([(k in config[T0_SECT]) for k in opt_names]):
        raise ValueError('incompatible T0 options: {}'.format(opt_names))
m_ref = config.getfloat(T0_SECT, 'M-ref', fallback=None)
if m_ref:
    tm_ref = m_ref * rd.qnms.T_MSUN
    logging.info('Reference mass: {} Msun ({} s)'.format(m_ref, tm_ref))
else:
    tm_ref = 1
t0ref = config.getfloat(T0_SECT, T0_KEYS['ref'], fallback=0)
if T0_KEYS['list'] in config[T0_SECT]:
    t0s = array(literal_eval(config.get(T0_SECT, T0_KEYS['list'])))
elif T0_KEYS['delta'] in config[T0_SECT]:
    dt0s = array(literal_eval(config.get(T0_SECT, T0_KEYS['delta'])))
    t0s = dt0s * tm_ref + t0ref
elif any([(k in config[T0_SECT]) for k in start_stop_step]):
    if not all([(k in config[T0_SECT]) for k in start_stop_step]):
        missing = [k for k in start_stop_step if k not in config[T0_SECT]]
        raise ValueError('missing start/stop/step options: {}'.format(missing))
    start, stop, step = [config.getfloat(T0_SECT, k) for k in start_stop_step]
    if start > 500 and t0ref > 100000000.0:
        logging.warning(
            'high reference time and stepping start---did you accidentally provide GPS times twice?'
            )
    t0s = arange(start, stop, step) * tm_ref + t0ref
else:
    raise ValueError(
        'no timing instructions in [{}] section; valid options are: {}'.
        format(T0_SECT, list(T0_KEYS.values())))
savetxt(PATHS['t0'], t0s)
if rerun and os.path.exists(PATHS['run_task']):
    os.remove(PATHS['run_task'])
TASK = (
    'mkdir -p {rundir}; cd {rundir}; ringdown_fit -o {result} {config} &> run.log\n'
    )
logging.info(f"Processing {len(t0s)} start times (recorded in {PATHS['t0']})")
for i, t0 in enumerate(t0s):
    cpath = PATHS['run_config'].format(i=i, t0=t0)
    rundir = os.path.dirname(cpath)
    os.makedirs(rundir, exist_ok=True)
    config_child = configparser.ConfigParser()
    config_child.read(PATHS['config'])
    config_child['target']['t0'] = str(t0)
    rpath = PATHS['run_result'].format(i=i, t0=t0)
    if os.path.exists(rpath):
        logging.info(f'Run {i} already exists. Will skip.')
    else:
        with open(cpath, 'w') as f:
            config_child.write(f)
        with open(PATHS['run_task'], 'a') as f:
            f.write(TASK.format(rundir=rundir, result=rpath, config=cpath))
logging.info('Done processing {} runs.'.format(len(t0s)))
NSLURM = 7 * args.nodes
NCPU = 4
if NSLURM * NCPU > 1280:
    w = 'Requested number of cores ({}) above 1280 user limit.'
    logging.warning(w.format(NSLURM * NCPU))
EXE = ['#! /usr/bin/env bash', 'cd {}'.format(outdir),
    'sbatch -p cca -n {} -c {} disBatch {}'.format(NSLURM, NCPU, PATHS[
    'run_task']), 'cd -']
epath = PATHS['exe']
with open(epath, 'w') as f:
    f.write('\n'.join(EXE))
st = os.stat(epath)
os.chmod(epath, st.st_mode | 73)
if args.submit:
    print('Submitting: {}'.format(epath))
    import subprocess
    subprocess.run(epath)
else:
    print('Submit by running: {}'.format(epath))

#!/usr/bin/env python
# coding: utf-8

# Copyright 2022
# Maximiliano Isi <max.isi@ligo.org>
#
# This program 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 2 of the License, or
# (at your option) any later version.
#
# This program 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 this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA 02110-1301, USA.

import sys
import os
from pylab import *
import ringdown as rd
import argparse
import configparser
import logging
import time
from ast import literal_eval


##############################################################################
# PARSE INPUT
##############################################################################

parser  = argparse.ArgumentParser(description="Analyze event ringdown at "
                                  "different times.")
parser.add_argument('config', help="Path to configuration file.")
parser.add_argument('-o', '--outdir', default=None, help="Output directory.")
parser.add_argument('--force', action='store_true', help="Run even if output "
                    "directory already exists (pre-existing runs will not be "
                    "repeated).")
parser.add_argument('-s', '--submit', action='store_true', help="Submit slurm "
                    "job automatically.")
parser.add_argument('--nodes', default=20, type=int,
                    help="Maximum number of nodes to request through SLURM.")
parser.add_argument('--seed', default=None, type=int, help="Random seed.")
parser.add_argument('-v', '--verbose', action='store_true')

args = parser.parse_args()

# set up logging and procure Git repo hash
if args.verbose:
    logging.getLogger().setLevel(logging.INFO)

try:
    import git
    repo = git.Repo(search_parent_directories=True)
    githash = repo.head.object.hexsha
    logging.info('GIT {}'.format(githash))
except ImportError:
    githash = 'unknown'
    logging.info("install gitpython to report repo hash")
except Exception:
    githash = 'unknown'
    logging.warning("failed to locate git repo")

# load config file 
config_path = os.path.abspath(args.config)
if not os.path.isfile(config_path):
    raise FileNotFoundError(f"unable to load: {config_path}")
logging.info(f"Loading config from {config_path}")
config = configparser.ConfigParser()
config.read(config_path)

# set random seed (purposedly fail if not provided)
seed = args.seed or config.getint('pipe', 'seed')
logging.info("Random seed set to {}".format(seed))
rng = np.random.default_rng(seed)

# determine run directory
outdir = args.outdir or config.get('pipe', 'outdir', fallback=None)
if not outdir:
    outdir = 'ringdown_pipe_{:.0f}'.format(time.time())
    logging.warning("No run dir provided, defaulting to {}".format(outdir))
outdir = os.path.abspath(outdir)
logging.info("Running in {}".format(outdir))

rerun = False
if os.path.exists(outdir):
    if args.force:
        logging.warning("Run directory already exists.")
        rerun = True
    else:
        raise FileExistsError("Run directory already exists. Exiting.")
else:
    os.makedirs(outdir)

PATHS = {
    'config': 'config.ini',
    'command': 'pipe.sh',
    'run_config': 'engine/{t0:.6f}/config.ini',
    'run_result': 'engine/{t0:.6f}/result.nc',
    'run_task': 'TaskFile',
    'exe': 'submit.sh',
    't0': 'start_times.txt',
}
PATHS = {k: os.path.join(outdir, v) for k,v in PATHS.items()}

# record config and arguments for reproducibility
with open(PATHS['config'], 'w') as f:
    config.write(f)

with open(PATHS['command'], 'w') as f:
    f.write("{}\n\n".format(' '.join(sys.argv)))
    for k,v in vars(args).items():
        f.write("# {}: {}\n".format(k,v))
    f.write("\n# {}".format(os.path.realpath(__file__)))
    f.write("\n# GIT hash: {}".format(githash))
    f.write("\n# ringdown v {}".format(getattr(rd, '__version__', 'unknown')))

# Identify target analysis times. There will be three possibilities:
#     1- listing the times explicitly
#     2- listing time differences with respect to a reference time
#     3- providing start, stop, step instructions to construct start times
#        (potentially relative to a reference time)
# Time steps/differences can be specified in seconds or M, if a reference mass
# is provided (in solar masses).

# Define valid options to specify the start times
T0_SECT = "pipe"
T0_KEYS = {
    'ref': 't0-ref',
    'list': 't0-list',
    'delta': 't0-delta-list',
    'step': 't0-step',
    'start': 't0-start',
    'stop': 't0-stop'
}
start_stop_step = [T0_KEYS[k] for k in ['start', 'stop', 'step']]

# First make sure that only compatible t0 options were provided
incompatible_sets = [['ref', 'list'], ['delta', 'list']]
incompatible_sets += [[k, 'delta'] for k in ['start', 'stop', 'step']]
for bad_set in incompatible_sets:
    opt_names = [T0_KEYS[k] for k in bad_set]
    if all([k in config[T0_SECT] for k in opt_names]):
        raise ValueError("incompatible T0 options: {}".format(opt_names))

# Look for a reference mass, to be used when stepping in time
m_ref = config.getfloat(T0_SECT, 'M-ref', fallback=None)
if m_ref:
    # reference time translating from solar masses
    tm_ref = m_ref * rd.qnms.T_MSUN
    logging.info("Reference mass: {} Msun ({} s)".format(m_ref, tm_ref))
else:
    # no reference mass provided, so will default to seconds
    tm_ref = 1

# Look for reference time to be used to construct start times
t0ref = config.getfloat(T0_SECT, T0_KEYS['ref'], fallback=0)

# Now we can safely interpret the options assuming one of three cases
if T0_KEYS['list'] in config[T0_SECT]:
    t0s = array(literal_eval(config.get(T0_SECT, T0_KEYS['list'])))
elif T0_KEYS['delta'] in config[T0_SECT]:
    dt0s = array(literal_eval(config.get(T0_SECT, T0_KEYS['delta'])))
    t0s = dt0s*tm_ref + t0ref
elif any([k in config[T0_SECT] for k in start_stop_step]):
    if not all([k in config[T0_SECT] for k in start_stop_step]):
        missing = [k for k in start_stop_step if k not in config[T0_SECT]]
        raise ValueError("missing start/stop/step options: {}".format(missing))
    # add a safety check here, in case the user mistakenly requests stepping
    # based on a GPS time and provides a reference GPS time
    start, stop, step  = [config.getfloat(T0_SECT, k) for k in start_stop_step]
    if start > 500 and t0ref > 1E8:
        logging.warning("high reference time and stepping start---did you "
                        "accidentally provide GPS times twice?")
    t0s = arange(start, stop, step)*tm_ref + t0ref
else:
    raise ValueError("no timing instructions in [{}] section; valid options "
                     "are: {}".format(T0_SECT, list(T0_KEYS.values())))

# Record start times
savetxt(PATHS["t0"], t0s)

##############################################################################
# SET UP RUNS
##############################################################################

if rerun and os.path.exists(PATHS['run_task']):
    os.remove(PATHS['run_task'])

TASK = "mkdir -p {rundir}; cd {rundir}; ringdown_fit -o {result} {config} "\
       "&> run.log\n"

logging.info(f"Processing {len(t0s)} start times (recorded in {PATHS['t0']})")
for i, t0 in enumerate(t0s):
    # Set up child configuration file
    cpath = PATHS['run_config'].format(i=i, t0=t0)
    rundir = os.path.dirname(cpath)
    os.makedirs(rundir, exist_ok=True)
    config_child = configparser.ConfigParser()
    config_child.read(PATHS['config'])

    # Set start time
    # Restore conditioning section
    config_child['target']['t0'] = str(t0)

    # Write out config file for this run
    rpath = PATHS['run_result'].format(i=i, t0=t0)
    if os.path.exists(rpath):
        logging.info(f"Run {i} already exists. Will skip.")
    else:
        with open(cpath, 'w') as f:
            config_child.write(f)
        with open(PATHS['run_task'], 'a') as f:
            f.write(TASK.format(rundir=rundir, result=rpath, config=cpath))

logging.info("Done processing {} runs.".format(len(t0s)))


##############################################################################
# SLURM WORKFLOW
##############################################################################

# rusty Broadwell nodes have 28 cores, so each can fit 7 4-core jobs at a time
# given that, make the number of jobs executed at a time a multiple of 7
NSLURM = 7*args.nodes
NCPU = 4

# a user may allocate up to 1280 cores at one time in a center partition, so
# issue a warning if this limit will be exceeded
if NSLURM*NCPU > 1280:
    w = "Requested number of cores ({}) above 1280 user limit."
    logging.warning(w.format(NSLURM*NCPU))

EXE = [
    '#! /usr/bin/env bash',
    'cd {}'.format(outdir),
    'sbatch -p cca -n {} -c {} disBatch {}'.format(NSLURM, NCPU, PATHS['run_task']),
    'cd -',
]

epath = PATHS['exe']
with open(epath, 'w') as f:
    f.write('\n'.join(EXE))
st = os.stat(epath)
os.chmod(epath, st.st_mode | 0o111)
if args.submit:
    print("Submitting: {}".format(epath))
    import subprocess
    subprocess.run(epath)
else:
    print("Submit by running: {}".format(epath))
