###############################################################################
# Class: JobSubmitter_sbatch_AWE
#
# Purpose:    Custom "sbatch" job submitter for AWE
# 
# Programmer: Brad Whitlock
# Date:       Tue Jan 15 15:29:23 GMT 2013
#
# Modifications:
#
# Satheesh Maheswaran, Paul Selby, Wed May 15 11:13:02 BST 2013
# Modified sbatch job submitter to work on AWE platforms
#
###############################################################################

class JobSubmitter_sbatch_AWE(JobSubmitter_sbatch):
    def __init__(self, launcher):
        super(JobSubmitter_sbatch_AWE, self).__init__(launcher)
        ppn = "8"
        if self.parallel.nn != None:
            ppn = int(math.ceil(float(self.parallel.np) / float(self.parallel.nn)))
            
    def TFileLoadModules(self, tfile):
        tfile.write("#SBATCH -J visit # Job Name\n")
        tfile.write("ulimit -l unlimited\n")
        tfile.write(". %s/init/bash\n" %GETENV("MODULESHOME"))
        tfile.write("module purge\n")
        tfile.write("export I_MPI_PLATFORM=auto\n")
        tfile.write("module load slurm intel/11.1\n")
        tfile.write("module load intelmpi/4.0.3-intel\n")
        
###############################################################################
# Class: JobSubmitter_mpirun_AWE
#
# Purpose:    Custom "mpirun" job submitter for AWE
# 
# Programmer: Satheesh Maheswaran, Paul Selby
# Date:       Tue Jan 15 15:29:23 GMT 2013
#
###############################################################################
 
class JobSubmitter_mpirun_AWE(JobSubmitter_mpirun):
    def __init__(self, launcher):
        super(JobSubmitter_mpirun_AWE, self).__init__(launcher)

    def Executable(self):
        return ["mpirun.awe"]
        
###############################################################################
# Class: JobSubmitter_qsub_AWE
#
# Purpose:    Custom "qsub" job submitter for AWE
# 
# Programmer: Satheesh Maheswaran, Paul Selby
# Date:       Tue Jan 15 15:29:23 GMT 2013
#
###############################################################################  
     
class JobSubmitter_qsub_AWE(JobSubmitter_qsub):
    def __init__(self, launcher):
        super(JobSubmitter_qsub_AWE, self).__init__(launcher)

        ncores = 16

        # Not 100% this should be in init method
        if self.parallel.nn == None:
          self.parallel.nn = int(math.ceil(float(self.parallel.np) / ncores))

        # Check ppn is in range
        ppn = int(self.PPN())
        
        if ppn > ncores:
            exit("Error: engine configured with %s on %s nodes (%sx)\n"
                 "Error: Number of procs must be no more than %sx number of nodes." %
                 (self.parallel.np, self.parallel.nn, ppn, ncores), 1)

    def mpirun(self):
        return ["mpirun.awe"]

    def TFileLoadModules(self, tfile):

        if self.launcher.IsRunningOnWillow():
            tfile.write("#MSUB -E       # Export MOAB environment variables\n")
            tfile.write("#MSUB -N visit # Job Name\n")
            tfile.write("#MSUB -A visit # Code name for accounting\n")
            tfile.write("#MSUB -j oe    # Joing stderr/stdout together\n")
            tfile.write(". %s/init/bash\n" %GETENV("MODULESHOME"))
            tfile.write("module purge\n")
            tfile.write("module load intel/12.1\n")
            tfile.write("module load bullxmpi/1.1.11-intel\n")
        elif self.launcher.IsRunningOnBlackthorn():
            tfile.write("#MSUB -E       # Export MOAB environment variables\n")
            tfile.write("#MSUB -N visit # Job Name\n")
            tfile.write("#MSUB -A visit # Code name for accounting\n")
            tfile.write("#MSUB -j oe    # Joing stderr/stdout together\n")
            tfile.write(". %s/init/bash\n" %GETENV("MODULESHOME"))
            tfile.write("module purge\n")
            tfile.write("module load intel/12.1\n")
            tfile.write("module load intelmpi/4.0.3-intel\n")
        return
        
###############################################################################
# Class: AWELauncher
#
# Purpose:    Custom launcher for AWE
# 
# Programmer: Satheesh Maheswaran, Paul Selby
# Date:       Tue Jan 15 15:29:23 GMT 2013
#
###############################################################################

class AWELauncher(MainLauncher):
    def __init__(self):
        super(AWELauncher, self).__init__()

    def IsRunningOnWillow(self):
        return self.sectorname() in ("willowa", "willowb")

    def IsRunningOnBlackthorn(self):
        return self.sectorname() in ("blackthorn")
    
    def IsRunningOnIvy(self):
        return self.sectorname() in ("ivy")

    def IsRunningOnBonsai(self):
        return self.sectorname() in ("bonsaia", "bonsaib")

    def IsRunningOnWorkstation(self):
        return GETENV("HPC_CLUSTER") == "LINUX_WORKSTATION"
        
    def Customize(self):
    
        ld_paths = self.splitpaths(GETENV("LD_LIBRARY_PATH"))
        paths = self.splitpaths(GETENV("PATH"))
        
        if self.IsRunningOnWillow():
            if GETENV("MODULESHOME") == "" :
                SETENV("MODULESHOME", "/awe/Modules")
                
        elif self.IsRunningOnBlackthorn() :
            if GETENV("MODULESHOME") == "" :
                SETENV("MODULESHOME", "/awe/Modules")
            if GETENV("MOABHOMEDIR") == "" :
                SETENV("MOABHOMEDIR", "/opt/moab/home")
                
            paths += ["/opt/moab/current/bin"]
        elif self.IsRunningOnIvy():
            if GETENV("MODULESHOME") == "" :
                SETENV("MODULESHOME", "/usr/share/modules")
                
            paths += ["/awe/slurm/bin"]
        elif self.IsRunningOnBonsai():
            paths += ["/opt/sgi/mpt/mpt-2.05/bin"]
            ld_paths += ["/opt/sgi/mpt/mpt-2.05/lib"]
        elif self.IsRunningOnWorkstation():
            paths += ["/awe/intel/impi/4.0.0/intel64/bin"]
        else:
            exit("Error: Unknown platform\n")
            
        SETENV("PATH", self.joinpaths(paths))
        SETENV("LD_LIBRARY_PATH", self.joinpaths(ld_paths))

    def JobSubmitterFactory(self, launch):
        if launch[:4] == 'msub':
            return JobSubmitter_qsub_AWE(self)
        elif launch[:6] == 'sbatch':
            return JobSubmitter_sbatch_AWE(self)
        elif launch[:6] == 'mpirun':
            if self.IsRunningOnBonsai() or self.IsRunningOnIvy() :
                return JobSubmitter_mpirun(self)
            else:
                return JobSubmitter_mpirun_AWE(self)
        return super(AWELauncher, self).JobSubmitterFactory(launch)

# Launcher creation function
def createlauncher():
    return AWELauncher()
