#!/usr/bin/env python3
import sh
import argparse

parser = argparse.ArgumentParser(
    description='Submit a job to SGE in a given python and shell environment.')
parser.add_argument('job_name', type=str, help='Name of the job.')
parser.add_argument('n_cpus', type=str, help='Number of cpus.')
parser.add_argument('cmds', type=str, help='Commands to execute.')
parser.add_argument('--venv', type=str, help='Virtual environment.',
                    default=None)
parser.add_argument('--worker', type=str, default=None,
                    help='Worker to execute the job (e.g. cpu08).')
parser.add_argument('--gpus', type=str, default=0,
                    help='Number of gpus.')
parser.add_argument('--mem', type=str, default="2G",
                    help='Amount of memory to allocate for the job.')
args = parser.parse_args()


qsub_args = [
    "-b", "y", "-cwd",
    "-N", args.job_name,
    "-l", "threads={}".format(args.n_cpus),
    "-l", "memory={}".format(args.mem),
]

if args.gpus > 0:
    qsub_args += ["-l", "gpu={}".format(args.gpus)]

if args.worker is not None:
    qsub_args += ["-l", "hostname={}".format(args.worker)]

envs = ['$HOME/deepmeth/deepmeth_analysis/env_vars.sh']

if args.venv is not None:
    envs.append("$HOME/{}/bin/activate".format(args.venv))

cmd = ";".join("source " + env for env in envs)
cmd += ";echo 'Executing:'"
cmd += ";echo '{}'".format(args.cmds)
cmd += ";echo ''"
cmd += ";time " + args.cmds
qsub_args.append(cmd)

print(sh.qsub(*qsub_args))
