#!/usr/bin/env python3
import subprocess
import uuid

try:
    import click
except ImportError:
    print("install click: `pip install click`")
    exit

PBS_TEMPLATE = r"""#!/usr/bin/env bash

#PBS -q {queue}
#PBS -m be
#PBS -k oe
#PBS -N {program_name}
#PBS -l nodes={nodes}:ppn={cores}
#PBS -l mem={mem}MB
#PBS -l walltime={walltime}:00:00


echo ------------------------------------------------------
echo -n 'Job is running on node '; cat $PBS_NODEFILE
echo ------------------------------------------------------
echo "PBS: originating queue is"  $PBS_O_QUEUE
echo "PBS: executing queue is"  $PBS_QUEUE
echo "PBS: working directory is"  $PBS_O_WORKDIR
echo "PBS: execution mode is"  $PBS_ENVIRONMENT
echo "PBS: job identifier is"  $PBS_JOBID
echo "PBS: job name is"  $PBS_JOBNAME
echo "PBS: node file is" $PBS_NODEFILE
echo "PBS: current home directory is" $PBS_O_HOME
echo "PBS: PATH =" $PBS_O_PATH
echo ------------------------------------------------------

module use --append $PBS_O_HOME/privatemodules

module load psi4/psi4-1.4a2-parallel_psi4numpy
module load mrcc/2020.1-intel-2019.0.045

export HOSTS_FILE=$PBS_NODEFILE

# create job tmpdir
TMPDIR=/scratch/$PBS_JOBID
mkdir -p $TMPDIR
cd $TMPDIR

# make psi scratch dir
mkdir -p psi4

# copy content of the pwd
cp -rf $PBS_O_WORKDIR/* $TMPDIR/

# set psi4 scratch dir
export PSI_SCRATCH="${{TMPDIR}}/psi4"
psi4 -n {cores} -i {program_name} | tee {program_name}.out

# remove psi4 scratch dir
rm -rf $TMPDIR/psi4

# copy back the results
cp -rf $TMPDIR/* $PBS_O_WORKDIR/
cd $PBS_O_WORKDIR

# remove job tmpdir
rm -rf $TMPDIR

# scrap the job number
jobid=$PBS_JOBID
NUMBER=${{jobid%.*}}

# move back job error & output
mv $PBS_O_HOME/*.o$NUMBER $PBS_O_WORKDIR/
mv $PBS_O_HOME/*.e$NUMBER $PBS_O_WORKDIR/
"""


def log_submition(nodes, cores, mem, walltime, queue, input):
    print("The job is being submitted with the following parameters:")
    print("")
    print(f"'queue'   : {queue}")
    print(f"'nodes'   : {nodes}")
    print(f"'cores'   : {cores}")
    print(f"'mem'     : {mem}")
    print(f"'walltime': {walltime}")


def render_template(*args, **kwargs):
    return PBS_TEMPLATE.format(**kwargs)


@click.command()
@click.option(
    "--nodes", default=1, help="number of requested nodes", type=click.IntRange(min=1)
)
@click.option(
    "--cores", default=1, help="number of cores per node", type=click.IntRange(min=1)
)
@click.option(
    "--mem",
    default=800,
    help="number of memory per node in MiB",
    type=click.IntRange(min=1),
)
@click.option(
    "--walltime",
    default=24,
    help="requested walltime in hours",
    type=click.IntRange(min=1),
)
@click.option(
    "--queue",
    default="maink",
    help="default submitting queue",
    type=click.Choice(["maink", "fat"]),
)
@click.argument("program_name")
def main(nodes, cores, mem, walltime, queue, program_name):
    """simple program to submit jobs to PBS queueing system"""
    log_submition(nodes, cores, mem, walltime, queue, program_name)
    batch_file_ = render_template(
        nodes=nodes,
        cores=cores,
        mem=mem,
        walltime=walltime,
        queue=queue,
        program_name=program_name,
    )
    uuid_ = uuid.uuid4()
    batch_file_name = f"pbs-job-{uuid_}.sh"
    with open(f"submition-{uuid_}.out", "w") as out_file:
        with open(batch_file_name, "w") as batch_file:
            batch_file.writelines(batch_file_)

        # try calling qsub with non-rerunable job
        # NOTE: qsub notes https://www.jlab.org/hpc/PBS/qsub.html
        subprocess.call(["qsub", "-r", "n", batch_file_name], stdout=out_file)


if __name__ == "__main__":
    main()