submit

#!/bin/python
from __future__ import print_function
import sys
import os
import argparse

class SlurmJob(object):
    def __init__(self,args):
        # Check input actually exists
        if not os.path.isfile(args.infile):
            sys.exit('The input file does not exist!')
        
        # Make sure we have the correct extension
        infile = args.infile.split('.')
        self.software = None
        try:
            extension = infile[1]
        except IndexError:
            sys.exit('Q-Chem input file must have .in or .inp extension')
        if (extension == 'in') or (extension == 'inp'):
            print('Extension: ',extension,'; assuming a Q-Chem job')
            self.software = 'qchem'
        elif (extension == 'com') or (extension == 'gjf'):
            print('Extension: ',extension,'; assuming a Gaussian 16 job')
            self.software = 'g16'
        else:
            sys.exit('Q-Chem input file must have .in or .inp extension \n \
                      Gaussian input file must have .com or .gjf extension \n \
                      ...please make the appropriate changes and resubmit.')
     
        self.name      = infile[0]
        self.infile    = args.infile 
        self.threads   = args.nthreads
        self.partition = args.partition
        self.time      = args.time
        self.script    = self.name+'.sh'

        # Test if the walltime is valid on the requested partitions 
        partition = self.partition.split(',')
        if (self.time > 672):
            sys.exit("Walltime of %s hours too long for Grace! \n Exiting..." % self.time)
        if (self.time > 168) and ('week' in partition):
            partition.remove('week')
        if (self.time > 24) and ('day' in partition):
            partition.remove('day')
        too_long = list(set(self.partition.split(',')).difference(partition))
        if too_long:
            self.partition = ','.join(partition) 
            print("Wall time of %s h too long for '%s' partition(s)." % (self.time, ','.join(too_long)))
            print("Modifying partition to: '%s'" % self.partition)


    def write_slurm_script(self):
        with open(self.script,'w') as f:
           f.write("#!/bin/bash\n")
           f.write("#SBATCH --job-name="+self.name+"\n")
           f.write("#SBATCH --nodes=1\n")
           f.write("#SBATCH --cpus-per-task="+str(self.threads)+"\n")
           f.write("#SBATCH --partition="+self.partition+"\n")
           if 'hammes_schiffer' in self.partition:
               f.write("#SBATCH -A hammes_schiffer\n")
           f.write("#SBATCH -t "+str(self.time).zfill(2)+":00:00\n")
           f.write("\n")
           if self.software == 'qchem':
               #f.write("module load qchem/5.1dev\n")
               #f.write("\n")
               #f.write(self.software+" -nt "+str(self.threads)+" "+self.infile+" "+self.name+".out\n")
               f.write(self.software+" -nt ${SLURM_CPUS_ON_NODE} -save "+self.infile+" "+self.name+".out ${SLURM_JOB_NAME}_${SLURM_JOB_ID}\n")
               f.write("\n")
               f.write("cp -r ${QCSCRATCH}/${SLURM_JOB_NAME}_${SLURM_JOB_ID}/AIMD ${SLURM_JOB_NAME}_${SLURM_JOB_ID}_AIMD")
               f.write("\n")
               f.write("rm slurm-${SLURM_JOB_ID}.out")
           elif self.software == 'g16':
               f.write("module load Gaussian/16-C.01_AVX\n")
               f.write("\n")
               f.write(self.software+" < "+self.infile+" > "+self.name+".log\n")
               f.write("\n")
               f.write("rm slurm-${SLURM_JOB_ID}.out\n")
               f.write("rm core.*")

        print(self.software+' submission script written to '+self.script) 
        print('To submit, do "sbatch '+self.script+'"') 

    
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Prepare Gaussian 16 or Q-Chem submit script on Grace.')
    parser.add_argument("infile", help="Gaussian 16 or Q-Chem input file")
    parser.add_argument("-nt", "--nthreads",metavar='N', help="CPUs-per-node",type=int,default=16)
    parser.add_argument("-p","--partition",metavar='P',help="partition",default="pi_hammes_schiffer,day")
    parser.add_argument("-t","--time",metavar='T',help="time (hours)",type=int,default=24)
    args = parser.parse_args()

    slurmjob = SlurmJob(args)
    slurmjob.write_slurm_script()