###############################################################################
# Class: JobSubmitter_qsub_LLNL
#
# Purpose:    Custom "qsub" job submitter for LLNL.
#
# Programmer: Brad Whitlock
# Date:       Thu May 17 14:22:04 PDT 2012
#
# Modifications:
#   Eric Brugger, Tue Oct 18 16:38:04 PDT 2016
#   I modified the script to set ppn when running on trinity.
#
#   Eric Brugger, Thu Dec 29 13:12:25 PST 2016
#   I modified the script to set ttc when running on agate.
#
#   Eric Brugger, Mon Jan  9 13:43:22 PST 2017
#   I modified the script to no longer set usettc when on inca.
#
#   Eric Brugger, Thu Mar 29 08:32:40 PDT 2018
#   I modified the script to set ttc when running on rztrona.
#
###############################################################################

class JobSubmitter_qsub_LLNL(JobSubmitter_qsub):
    def __init__(self, launcher):
        super(JobSubmitter_qsub_LLNL, self).__init__(launcher)
        self.usettc = self.sectorname() in ("agate", "borax", "rzalastor", "rztrona")
        self.useppn = (self.sectorname() == "graph") or (self.sectorname() == "tr-fe")

###############################################################################
# Class: JobSubmitter_mpirun_LLNL
#
# Purpose:    Custom "mpirun" job submitter for LLNL.
#
# Programmer: Eric Brugger
# Date:       Thu Feb 15 09:38:48 PST 2018
#
###############################################################################

class JobSubmitter_mpirun_LLNL(JobSubmitter_mpirun):
    def __init__(self, launcher):
        super(JobSubmitter_mpirun_LLNL, self).__init__(launcher)

    def Executable(self):
        if self.sectorname() in ("quartz", "borax"):
            return ["/usr/tce/packages/openmpi/openmpi-2.1.0-gcc-4.9.3/bin/mpirun"]
        return super(JobSubmitter_mpirun_LLNL, self).Executable()

###############################################################################
# Class: JobSubmitter_bsub_LLNL
#
# Purpose:    Custom "bsub" job submitter for LLNL.
#
# Programmer: Eric Brugger
# Date:       Fri Nov 30 14:15:23 PST 2018
#
###############################################################################

class JobSubmitter_bsub_LLNL(JobSubmitter):
    def __init__(self, launcher):
        super(JobSubmitter_bsub_LLNL, self).__init__(launcher)

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

    def CreateCommand(self, args, debugger):
        bsub, sublauncher = self.LauncherAndSubLauncher()
        parcmd = self.Executable()
        parcmd = parcmd + ["-csm", "y"]
        if self.parallel.launchargs != None:
            parcmd = parcmd + self.parallel.launchargs
        if self.parallel.np != None:
            parcmd = parcmd + ["-n", str(int(self.parallel.np)+1)]
        if self.parallel.partition != None:
            parcmd = parcmd + ["-q", self.parallel.partition]
        if self.parallel.time != None:
            parcmd = parcmd + ["-W", self.parallel.time]
        if self.parallel.bank != None:
            parcmd = parcmd + ["-G", self.parallel.bank]
        nodes = self.parallel.nn
        if nodes == None:
            nodes = self.parallel.np
        ppn = str(int(math.ceil(float(self.parallel.np) / float(nodes))))
        parcmd = parcmd + ["-R", "1*{select[(LN)&&(type==any)] order[r15s:pg] span[hosts=1]}+%s*{select[(CN)&&(type==any)] order[r15s:pg] span[ptile=%s]}" % (self.parallel.np,ppn)]
        parcmd = parcmd + [sublauncher]
        if self.parallel.sublaunchargs != None:
            parcmd = parcmd + self.parallel.sublaunchargs
        parcmd = parcmd + self.VisItExecutable()
        parcmd = parcmd + ["-plugindir", GETENV("VISITPLUGINDIR")]
        parcmd = parcmd + ["-visithome", GETENV("VISITHOME")]
        parcmd = parcmd + ["-visitarchhome", GETENV("VISITARCHHOME")]
        parcmd = parcmd + args
        parcmd = debugger.CreateCommand(parcmd)
        return parcmd

###############################################################################
# Class: LLNLLauncher
#
# Purpose:    Custom launcher for LLNL
#
# Programmer: Brad Whitlock
# Date:       Thu May 17 14:22:04 PDT 2012
#
# Modifications:
#   Eric Brugger, Mon Nov 19 12:04:57 PST 2012
#   I modified the script to add some paths to the LD_LIBRARY_PATH if running
#   on cielo.
#
#   Eric Brugger, Tue Nov 20 11:25:07 PST 2012
#   I removed the DetermineArchitecture method, which launched a chaos 4
#   version on some machines. Those machines are either retired or have been
#   upgraded to chaos 5.
#
#   Brad Whitlock, Wed Jun  5 16:24:26 PDT 2013
#   Special handling of networking arguments for BG/Q.
#
#   Cyrus Harrison, Tue Aug 23 08:21:12 PDT 2016
#   Added check for sequioa (seqlac)
#
#   Eric Brugger, Thu Jan 19 14:51:03 PST 2017
#   I added code to map login node names on jade and agate to their IP
#   addresses since they aren't in DNS.
#
#   Eric Brugger, Wed Mar 22 13:43:49 PDT 2017
#   I added the MPI library paths to the LD_LIBRARY_PATH for borax, quartz,
#   rzgenie and rztrona.
#
#   Eric Brugger, Thu Jun 15 11:44:37 PDT 2017
#   I added the MPI library paths to the LD_LIBRARY_PATH for rztopaz.
#
#   Eric Brugger, Wed Mar 22 14:17:24 PDT 2017
#   I added the MPI library paths to the LD_LIBRARY_PATH for jade and agate.
#   I removed the logic that sets the LD_LIBRARY_PATH if running on cielo.
#
#   Eric Brugger, Mon Mar 19 12:55:14 PDT 2018
#   I set visitarch to linux-x86_64-toss3 for all the toss3 machines.
#
#   Eric Brugger, Mon Apr 30 09:18:14 PDT 2018
#   I added rzhasgpu to the list of toss3 machines.
#
#   Eric Brugger, Fri May 25 11:55:08 PDT 2018
#   I added pascal to the list of toss3 machines.
#
#   Cyrus Harrison, Wed Jun 27 12:41:03 PDT 2018
#   I added catalyst to the list of toss3 machines.
#
#   Eric Brugger, Thu Jul 26 13:48:28 PDT 2018
#   I added surface to the list of toss3 machines.
#
#   Eric Brugger, Fri Aug 17 16:04:16 PDT 2018
#   I added syrah to the list of toss3 machines.
#
#   Eric Brugger, Thu Aug 30 12:27:37 PDT 2018
#   I added max to the list of toss3 machines.
#
#   Eric Brugger, Thu Nov  8 15:44:54 PST 2018
#   I added rzalastor and zin to the list of toss3 machines.
#
#   Eric Brugger, Fri Nov 30 14:15:23 PST 2018
#   I modified the launcher to use a custom launcher on lassen and sierra.
#
#   Eric Brugger, Wed Dec 12 11:54:00 PST 2018
#   I added pinot, rzansel and rzslic to the list of toss3 machines.
#
#   Eric Brugger, Tue Dec 18 14:32:09 PST 2018
#   I added czvnc and rzvnc to the list of toss3 machines.
#
###############################################################################

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

    # Get the IP address for eth0.
    def GetIPAddress(self):
        p = subprocess.Popen(["/sbin/ifconfig"], stdin=subprocess.PIPE, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
        output = p.communicate()
        iplines = [x for x in string.split(output[0], "\n") if string.find(x, "inet addr") != -1]
        start = string.find(iplines[0], "inet addr")
        end = string.find(iplines[0], "Bcast:")
        IP = string.replace(iplines[0][start + 10:end], " ", "")
        return IP

    def Customize(self):
        # ----
        # Hoth @ LLNL
        # ----
        if self.nodename() == "hoth" and self.domainname() == "llnl.gov":
            SETENV("LIBGL_ALWAYS_INDIRECT", "1")

        #
        # BG/Q networking changes for parallel engine. We override the host
        # with an IP address.
        #
        if self.parallelArgs.parallel and \
           (string.find(self.generalArgs.exe_name, "_par") != -1 or \
            string.find(self.generalArgs.exe_name, "_ser") != -1):
            if self.sectorname() == "vulcanlac" or \
               self.sectorname() == "rzuseqlac" or \
               self.sectorname() == "seqlac":
                self.generalArgs.host = self.GetIPAddress()
                self.generalArgs.guesshost = 0
                self.generalArgs.sshtunneling = 0
                self.generalArgs.noloopback = 1

        #
        # Convert the host name to the ip address on jade.
        #
        if self.generalArgs.host == "jade188.llnl.gov":
            self.generalArgs.host = "130.106.201.19"
        elif self.generalArgs.host == "jade380.llnl.gov":
            self.generalArgs.host = "130.106.201.20"
        elif self.generalArgs.host == "jade386.llnl.gov":
            self.generalArgs.host = "130.106.201.57"
        elif self.generalArgs.host == "jade764.llnl.gov":
            self.generalArgs.host = "130.106.201.58"
        elif self.generalArgs.host == "jade770.llnl.gov":
            self.generalArgs.host = "130.106.201.59"
        elif self.generalArgs.host == "jade962.llnl.gov":
            self.generalArgs.host = "130.106.201.60"
        elif self.generalArgs.host == "agate2.llnl.gov":
            self.generalArgs.host = "130.106.204.1"
        elif self.generalArgs.host == "agate5.llnl.gov":
            self.generalArgs.host = "130.106.204.3"

        #
        # Set the LD_LIBRARY_PATH to include the path to MPI on borax,
        # catalyst, jade, mica, pascal, quartz, rzgenie, rzhasgpu, rztopaz,
        # rztrona, surface, syrah and agate.
        #
        if self.sectorname() == "borax" or \
            self.sectorname() == "catalyst" or \
            self.sectorname() == "czvnc" or \
            self.sectorname() == "jade" or \
            self.sectorname() == "max" or \
            self.sectorname() == "mica" or \
            self.sectorname() == "pascal" or \
            self.sectorname() == "pinot" or \
            self.sectorname() == "quartz" or \
            self.sectorname() == "rzalastor" or \
            self.sectorname() == "rzansel" or \
            self.sectorname() == "rzgenie" or \
            self.sectorname() == "rzhasgpu" or \
            self.sectorname() == "rzslic" or \
            self.sectorname() == "rztopaz" or \
            self.sectorname() == "rztrona" or \
            self.sectorname() == "rzvnc" or \
            self.sectorname() == "surface" or \
            self.sectorname() == "syrah" or \
            self.sectorname() == "zin" or \
            self.sectorname() == "agate":
            mpi_ld_library_paths = ["/usr/tce/packages/mvapich2/mvapich2-2.2-intel-16.0.3/lib", "/usr/tce/packages/intel/intel-16.0.3/lib/intel64"]
            SETENV("LD_LIBRARY_PATH", self.joinpaths(mpi_ld_library_paths))
            self.visitarch = "linux-x86_64-toss3"

    #
    # Override the JobSubmitterFactory method so the custom job submitter can
    # be returned.
    #
    def JobSubmitterFactory(self, launch):
        if launch[:4] == "qsub" or launch[:4] == "msub":
            return JobSubmitter_qsub_LLNL(self)
        if launch[:6] == "mpirun":
            return JobSubmitter_mpirun_LLNL(self)
        if launch[:4] == "bsub":
            if self.sectorname() == "lassen" or \
                self.sectorname() == "sierra":
                return JobSubmitter_bsub_LLNL(self)
        return super(LLNLLauncher, self).JobSubmitterFactory(launch)

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