#!/usr/bin/env python
import datetime
import os
import sys
import signal

# from http://pguides.net/python/timeout-a-function
# June 2011
class TimeoutException(Exception):
    pass

def timeout(timeout_time, default):
    def timeout_function(f):
        def f2(*args):
            def timeout_handler(signum, frame):
                raise TimeoutException()
            
            old_handler = signal.signal(signal.SIGALRM, timeout_handler)
            signal.alarm(timeout_time) # trigger alarm in timeout_time seconds
            try:
                retval = f()
            except TimeoutException:
                return default
            finally:
                signal.signal(signal.SIGALRM, old_handler)
            signal.alarm(0)
            return retval
        return f2
    return timeout_function
    
# end from http://pguides.net/python/timeout-a-function
# the @ is a decorator
@timeout(10, "None")
def get_batchsystem():
    print "Please enter batch system: ",
    name = sys.stdin.readline()
    return name.rstrip("\n")

def setup_batch(forced=False):
    # first time use on this machine and userid
    defaultfile = os.path.expanduser('~')+"/.submit__defaults"
    if forced or not os.path.isfile(defaultfile):
        defaults = open(defaultfile,"w")
        print "What batch system do you use?"
        print "Options are"
        print "M   moab/torque"
        print "S   sun/oracle gridengine"
        print "N   none"
        batchflag = get_batchsystem()
        batchflag = batchflag.upper()[0]
        defaults.write(batchflag)
        defaults.close()
    else:
        defaults = open(defaultfile,"r")
        batchflag = defaults.readline()
        batchflag = batchflag.upper()[0]
        defaults.close()
    if batchflag!='M' and batchflag != 'S':
        batchflag = 'N'
    return batchflag



def sge(program,parmfile,logfile,nodes,wallhour,dobackfill):
    submitfile = "job"+logfile+".sub"
    job = open(submitfile,"w")
    job.write("#!/bin/bash\n")
    job.write("#$ -m e\n")
    job.write("#$ -j y\n")
    job.write("#$ -S /bin/bash\n")
    job.write("#$ -l h_cpu="+wallhour+" \n")
    job.write("#$ -V\n")
    job.write("#$ -cwd\n")
    if nodes > 1:
        job.write("mpirun -np "+str(nodes)+" "+program+" "+parmfile+" -nomenu\n") 
        job.close()
        print  "qsub -pe mpi ",nodes, submitfile, "(",program, parmfile, wallhour, ")";
        os.system("qsub -pe mpi "+str(nodes)+" "+submitfile)
    else:
        job.write(program+" "+parmfile+" -nomenu\n")
        job.close()
        print  "qsub ",submitfile, "(",program, parmfile, wallhour, ")";
        os.system("qsub "+submitfile)

def moab(program,parmfile,logfile,nodes,wallhour,dobackfill):
    submitfile = logfile+".sub"
    job = open(submitfile,"w")
    job.write("#!/bin/bash\n")
    job.write("#MOAB -o "+logfile+"\n")
    job.write("#MOAB -j oe\n")
    if nodes!=1:
        job.write("#MOAB -l nodes="+str(nodes)+"\n")
        job.write("#MOAB -l walltime="+wallhour+"\n")
        job.write("cd $PBS_O_WORKDIR\n")
        #job.write("mpirun -mca btl ^sm "+program+" "+parmfile+"-nomenu\n")
        job.write("mpirun "+program+" "+parmfile+"-nomenu\n")
    else:
        job.write("#MOAB -l walltime="+wallhour+"\n")
        job.write("cd $PBS_O_WORKDIR\n")
        job.write(program+" "+parmfile+" -nomenu\n")
    job.close()
    if dobackfill:
        print  "msub -q backfill ", submitfile, "(",program, parmfile, wallhour,")";
        os.system("msub -q backfill "+submitfile)
    else:
        print  "msub ", submitfile, "(",program, parmfile, wallhour,")";
        os.system("msub "+submitfile)

     
def standard(program,parmfile,logfile,nodes,wallhour,dobackfill):
    if nodes >1:
        os.system("mpirun -np "+str(nodes)+" "+program+" "+parmfile+" -nomenu")
    else:
        os.system(program+" "+parmfile+" -nomenu")
                
def wall(wallhour):    
    (hour,minute,second) =  [int(x) for x in wallhour.split(":")]
    if hour*3600 + minute*60 + second < 4*3600:
        backfillqueue = True
    else:
        backfillqueue = False
    return backfillqueue


def submit():
    today = str(datetime.date.today())
    definitions = ["","migrate-n-mpi","parmfile", "job"+today, 2,"4:00:00"]
    myhelp = set(["--help", "--h", "-h", "-help"])
    setup = set(["--setup"])
    arguments = sys.argv
    argc = len(arguments)
    definitions[:argc] = arguments[1:]
    print definitions
    
    needhelp = len(myhelp.intersection(definitions)) > 0
    
    if needhelp:
        print "Syntax:\n",
        print "submit --help   prints this help"
        print "submit --setup  sets up the batch queue"
        print "                default [default=none, moab/torque, sun grid engine]"
        print "submit program argument1 jobname numberCPU wallclockTime\n",
        print "Explanation [default in brackets]:\n",
        print "program         the executable needs to be an MPI program [migrate-n]\n",
        print "                if nodes =1 then assumed not MPI\n",
        print "argument1       one single argument, if no argument use \" \" [parmfile]\n",
        print "jobname         name of the job [job]\n",
        print "numberCPU       number of nodes used to run the job [5]\n",
        print "wallclockTime   Allowd runtime in hours (use only full hours) [4]"
        sys.exit()

    # forcing a setup of the batch system
    needsetup = len(setup.intersection(definitions)) > 0
    if needsetup:
        setup_batch(True)
        print "Submit is now setup with the correct batchqueue system, please rerun toe command"
        sys.exit()
        
    batchflag = setup_batch()
    
    (program, parmfile, logfile, nodes, wallhour) = definitions[:5]
    dobackfill = wall(wallhour)

    if batchflag == 'S':
        sge(program,parmfile,logfile,nodes,wallhour,dobackfill)
    if batchflag == 'M':
        moab(program,parmfile,logfile,nodes,wallhour,dobackfill)
    if batchflag == 'N':
        standard(program,parmfile,logfile,nodes,wallhour,dobackfill)
    sys.exit()



if __name__ == '__main__':
	submit()


