#!/usr/bin/env python

import os
import argparse
    

parser = argparse.ArgumentParser(description='Baseline Run')
parser.add_argument('--CODE_PATH', type=str, default="./",
                    help='Directory of file.')
parser.add_argument('--Cmd_file', type=str, default="main.py",
                    help='exact python file we need to execute')
parser.add_argument('--OutputDir', type=str, default="slurmRun",
                    help='exact python file we need to execute')
parser.add_argument('--Cmd_args', type=str, default="",
                    help='arguments for cmd')
args = parser.parse_args()
SlurmPath=os.path.join(args.CODE_PATH,args.OutputDir)
os.makedirs(SlurmPath, exist_ok = True)
job_file=os.path.join(SlurmPath,"run.sh")
output_file=os.path.join(SlurmPath,"run")

with open(job_file,"w") as fh:
    fh.writelines("#!/bin/bash\n")
    fh.writelines("#SBATCH --job-name=%s.job\n" % args.Cmd_file)
    fh.writelines("#SBATCH --output=%s.out\n" % output_file)
    fh.writelines("#SBATCH --error=%s.err\n" % output_file)
    fh.write('#SBATCH --gres=gpu:1\n')
    fh.write('#SBATCH --partition=titan-giant    # Partition to submit to \n')
    #fh.writelines("#SBATCH --time=2-00:00\n")
    fh.writelines("#SBATCH --mem=12000\n")
    fh.writelines("#SBATCH --qos=normal\n")
    fh.writelines("#SBATCH --mail-type=ALL\n")
    fh.writelines("#SBATCH --mail-user=$USER@cs.utah.edu\n")
    model_command = ' '.join([
                "python",
                os.path.join(args.CODE_PATH,args.Cmd_file),args.Cmd_args
        ])
    fh.write(model_command+"\n")
    fh.write('exit\n')
os.system("sbatch %s" %job_file)