#!/usr/bin/env python

import os
import stat
import uuid
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--dag-file', required=True)
parser.add_argument('--modes', required=True, nargs='+')
parser.add_argument('--seed', required=True, type=int)
parser.add_argument('injfiles', nargs='+')

opts = parser.parse_args()

seed = opts.seed

configdir = "/work/cdcapano/projects/gated_gaussian/gw190521/BH-Spectroscopy-GW190521/configuration/kerr"

cmdtmplt = """
pycbc_inference \\
    --nprocesses 8 \\
    --config-files {configdir}/{modes}/KERR-{modes}-{time}MS.ini \\
    --config-overrides data:injection-file:{injfile} \\
        data:trigger-time:FROM_INJECTION:tref \\
        static_params:tref:FROM_INJECTION:tref \\
        sampler:checkpoint_time_interval:7200 \\
        model:name:gated_gaussian_margpol \\
    --config-delete \\
        variable_params:polarization \\
        prior-polarization \\
    --output-file {output}.hdf \\
    --seed {seed} \\
    --verbose &> {output}.log
"""

subtmplt = """universe = vanilla
executable = {runfile}
initialdir = {rundir}
notification = NEVER
log = condor_log-{modes}-{time}
output = condor_out-{modes}-{time}
error = condor_error-{modes}-{time}
getenv = True
request_cpus = 8
accounting_group = cbc.imp.pe
request_memory = 40G
queue
"""

subfiles = []
for injfile in opts.injfiles:
    ii = int(os.path.basename(injfile).replace('inj', '').replace('.hdf', ''))
    rundir = os.path.basename(injfile).replace('.hdf', '')
    injfile = os.path.abspath(injfile)
    try:
        os.mkdir(rundir)
    except FileExistsError:
        pass
    for time in range(-9, 27, 3):
        if time < 0:
            time = 'M{}'.format(str(abs(time)).zfill(2))
        else:
            time = str(time).zfill(2)
        for modes in opts.modes:
            runsh = 'run{}-{}-{}.sh'.format(ii, modes, time)
            runfn = '/'.join([rundir, runsh])
            with open(runfn, 'w') as fp:
                print('#!/usr/bin/bash', file=fp)
                print('set -e', file=fp)
                output = 'samples-{}-{}'.format(modes, time)
                cmd = cmdtmplt.format(configdir=configdir, output=output,
                                      modes=modes, time=time,
                                      injfile=injfile, seed=seed)
                print(cmd, file=fp)
            seed += 78236
            st = os.stat(runfn)
            os.chmod(runfn, st.st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
            subfn = '{}/run{}-{}-{}.sub'.format(rundir, ii, modes, time)
            with open(subfn, 'w') as fp:
                sub = subtmplt.format(runfile=runfn,
                                      rundir=os.path.abspath(rundir),
                                      modes=modes, time=time)
                print(sub, file=fp)
            subfiles.append(subfn)

dagtmplt = "JOB {jobid} {subfn}"
with open(opts.dag_file, 'w') as fp:
    for subfn in subfiles:
        dagjob = dagtmplt.format(jobid=str(uuid.uuid4()).replace('-', ''),
                                 subfn=subfn) 
        print(dagjob, file=fp)