###############################################################################
# Class: JobSubmitter_ibrun_TACC
#
# Purpose:    Custom "ibrun" job submitter for TACC.
#
# Programmer: Brad Whitlock
# Date:       Thu May 17 14:22:04 PDT 2012
#
# Modifications:
#
#   Paul Navratil, Wed Jan  9 12:03:43 CST 2013
#   Removed deprecated Ranger and Spur configurations
#
###############################################################################

class JobSubmitter_ibrun_TACC(JobSubmitter_ibrun):
    def __init__(self, launcher):
        super(JobSubmitter_ibrun_TACC, self).__init__(launcher)

    def Executable(self):
        return ["ibrun", "tacc_xrun", "tacc_affinity"]

###############################################################################
# Class: JobSubmitter_qsub_TACC
#
# Purpose:    Custom "qsub" job submitter for TACC.
#
# Programmer: Brad Whitlock
# Date:       Thu May 17 14:22:04 PDT 2012
#
# Modifications:
#
#   Hank Childs, Fri Oct 26 11:06:18 PDT 2012
#   Remove the "as" part of the exception, as it isn't compliant with Python
#   2.6.
#
#   Paul Navratil, Wed Jan  9 12:03:43 CST 2013
#   Removed deprecated Ranger and Spur configurations
#
###############################################################################

class JobSubmitter_qsub_TACC(JobSubmitter_qsub):
    def __init__(self, launcher):
        super(JobSubmitter_qsub_TACC, self).__init__(launcher)

    def Executable(self):
        return ["qsub"]

    def AddEnvironment(self):
        args = []
        if not self.IsRunningOnRanger():
            args = super(JobSubmitter_qsub_TACC, self).AddEnvironment()
        return args

    def SetupPPN(self, nodes, procs, ppn, use_vis):
        thread_factor = 16
        if self.IsRunningOnLonghorn():
            thread_factor = 8
        if self.IsRunningOnLonestar():
            thread_factor = 12
        if self.parallel.nn != None:
            tacc_way = ppn
            tacc_procs = str(thread_factor * int(nnodes))
        else:
            # if nodes not set explicitly assume 1way
            tacc_way = "1"
            tacc_procs = str(thread_factor * int(procs))
        args = ["-pe", "%sway" % tacc_way, tacc_procs] 
        return args

    def AdditionalSetup(self):
        if self.IsRunningOnLonghorn():
            return ["-P", "vis"]
        return []

###############################################################################
# Class: TACCLauncher
#
# Purpose:    Custom launcher for TACC
#
# Programmer: Brad Whitlock
# Date:       Thu May 17 14:22:04 PDT 2012
#
# Modifications:
#
###############################################################################

class TACCLauncher(MainLauncher):
    def __init__(self):
        super(TACCLauncher, self).__init__()
        self.longhorn = 1
        self.lonestar = 0
        self.stampede = 0

    def IsRunningOnLonghorn(self):
        return self.longhorn

    def IsRunningOnLonestar(self):
        return self.lonestar

    def IsRunningOnStampede(self):
        return self.stampede

    def Customize(self):
        if self.parallelArgs.parallel and  \
           self.generalArgs.exe_name == "engine":
            if "longhorn" in self.hostname():
                self.longhorn = 1
            if "ls4" in self.hostname():
                self.lonestar = 1
            if "stampede" in self.hostname():
                self.stampede = 1
            msg = """

Attempting to launch parallel engine at TACC.

For more information about TACC resources and configurations, see:";
\thttp://www.tacc.utexas.edu/user-services/user-guides
"""
            self.message(msg)

    #
    # Override the JobSubmitterFactory method so the custom job submitter can
    # be returned.
    #
    def JobSubmitterFactory(self, launch):
        if launch == "ibrun":
            return JobSubmitter_ibrun_TACC(self)
        elif launch[:4] == "qsub":
            return JobSubmitter_qsub_TACC(self)
        return super(TACCLauncher, self).JobSubmitterFactory(launch)

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