import math

###############################################################################
# Class: JobSubmitter_aprun_NERSC
#
# Purpose:    Custom "aprun" job submitter for NERSC.
#
# Programmer: Brad Whitlock
# Date:       Thu May 17 14:22:04 PDT 2012
#
# Modifications:
#    Gunther H. Weber, Fri Oct  5 16:49:04 PDT 2012
#    Use fully qualified path name for aprun
#
###############################################################################

class JobSubmitter_aprun_NERSC(JobSubmitter_aprun):
    def __init__(self, launcher):
        super(JobSubmitter_aprun_NERSC, self).__init__(launcher)

    #
    # Override the name of the aprun executable and use fully qualified path
    # (Currently the path is identical on Edison and Hopper.)
    #
    def Executable(self):
        if self.launcher.nersc_host == "edison":
            return ["env", "DISPLAY=", "CRAY_ROOTFS=DSL", "/opt/cray/alps/default/bin/aprun"]
        else:
            return ["env", "DISPLAY=", "CRAY_ROOTFS=DSL", "/usr/common/usg/altd/1.0/bin/aprun"]

###############################################################################
# Class: JobSubmitter_mpirun_NERSC
#
# Purpose:    Custom "mpirun" job submitter for NERSC.
#
# Programmer: Gunther H. Weber
# Date:       Fri Oct  5 16:48:52 PDT 2012
#
# Modifications:
#
###############################################################################

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

    #
    # Use fully qualified path to mpirun executable
    #
    def Executable(self):
        # Currently, this submitter will only be called on Euclid
        return ["/usr/common/usg/openmpi/1.4.4/gnu/bin/mpirun"]

###############################################################################
# Class: JobSubmitter_qsub_NERSC
#
# Purpose:    Custom "qsub" job submitter for NERSC.
#
# Programmer: Brad Whitlock
# Date:       Thu May 17 14:22:04 PDT 2012
#
# Modifications:
#
###############################################################################

class JobSubmitter_qsub_NERSC(JobSubmitter_qsub):
    def __init__(self, launcher):
        super(JobSubmitter_qsub_NERSC, self).__init__(launcher)
        # FIXME: HACK: Set up VisIt socket relay on MOM node for parallel engine by overriding settings in launcher class
        launcher.loginnodeport = launcher.generalArgs.port
        launcher.generalArgs.port = "$MOM_PORT"
        launcher.loginnodehost = launcher.generalArgs.host
        launcher.generalArgs.host = "$MOM_HOST"
        launcher.generalArgs.noloopback = 1

    def CreateFilename(self, root):
        if  self.launcher.nersc_host == "edison" or self.launcher.nersc_host == "grace" or self.launcher.nersc_host == "hopper":
            tdate = time.asctime()[11:19]
            tuser = self.launcher.username()
            return os.path.join(GETENV("HOME"), "%s.%s.%s" % (root, tuser, tdate))
        return super(JobSubmitter_qsub_NERSC, self).CreateFilename(root)

    def Executable(self):
        if self.launcher.nersc_host == "carver":
            return ["/usr/syscom/opt/torque/default/bin/qsub"]
        elif self.launcher.nersc_host == "edison":
            return ["/opt/torque/default/bin/qsub"]
        elif self.launcher.nersc_host == "grace":
            return ["/usr/common/nsg/bin/qsub"]
        elif self.launcher.nersc_host == "hopper":
            return ["/opt/torque/default/bin/qsub"]
        else:
            print "Error: unknown NERSC host '%s'. Using default qsub executable." % self.launcher.nersc_host
            print "If this does not work, please contact consult@nersc.gov."
            return super(JobSubmitter_qsub_NERSC, self).Executable(self)

    def aprun(self):
        # Currently, this method will only be called on Edison or Hopper.
        # (Currently the path to aprun is the same on Edison and Hopper.)
        if self.launcher.nersc_host == "edison":
            return ["env", "CRAY_ROOTFS=DSL", "/opt/cray/alps/default/bin/aprun"]
        else:
            return ["env", "CRAY_ROOTFS=DSL", "/usr/common/usg/altd/1.0/bin/aprun"]

    def mpirun(self):
        # Currently, this method will only be called on Carver.
        return ["/usr/common/usg/openmpi/1.4.5/gcc/bin/mpirun"]
   
    def TFileSetup(self, tfile):
        tfile.write("cd %s\n" % os.path.abspath(os.curdir))
        tfile.write("ulimit -c 0\n")
        if self.launcher.nersc_host == "edison" or self.launcher.nersc_host == "grace" or self.launcher.nersc_host == "hopper":
            tfile.write("MOM_HOST=$(hostname)\n")
            relay = os.path.join(self.launcher.visitbindir, "visit_socket_relay")
            cmd = "%s %s %s > %s.port\n" % (relay, self.launcher.loginnodehost, self.launcher.loginnodeport, self.tfilename)
            tfile.write(cmd)
            tfile.write("MOM_PORT=$(cat %s.port)\n" % self.tfilename)
            tfile.write("rm %s.port\n" % self.tfilename)
            tfile.write("eval $(modulecmd sh unload xt-shmem)\n")

    def SetupPPN(self, nodes, procs, ppn, use_vis):
        if self.launcher.nersc_host == "edison" or self.launcher.nersc_host == "grace" or self.launcher.nersc_host == "hopper":
            args = []
            if self.parallel.workingdir != None:
                args = args + ["-d", self.parallel.workingdir]
            if self.parallel.nn != None:
                if self.launcher.nersc_host == "edison":
                    nprocpernode = 16
                else:
                    nprocpernode = 24
                nwidth = nprocpernode * int(nodes)
                # If the number of nodes is set, we need to modify mppwidth to ensure
                # allocation of the proper number of nodes. The number of nodes allocated
                # is mppwidth divided by the number of cores per node (compare
                # http://www.nersc.gov/users/computational-systems/hopper/running-jobs/batch-jobs/).
                args = args + ["-l", "mppwidth=%s" % str(nwidth)]
            else:
                # Number of nodes is not specified. We use all cores on each node and do not
                # need to modify mppdith to ensure the allocation of the appropriate number of
                # nodes.
                args = args + ["-l", "mppwidth=%s" % procs]
        elif self.launcher.nersc_host == "carver":
            new_nodes = nodes
            new_ppn = ppn
            if self.parallel.partition == "reg_xlmem":
                # There is a limit of one node for reg_xlmem
                new_nodes = 1
                new_ppn = procs
            else:
                # Even though Carver nodes may have a different number
                # of cores per nodes, this is a reasonable default for most
                # nodes. It is better to set nodes and cores manually in the
                # launch profile.
                if nodes == None:
                    new_nodes = math.ceil(procs / 8)
                    new_ppn = 8
            args = super(JobSubmitter_qsub_NERSC, self).SetupPPN(new_nodes, procs, new_ppn, use_vis)
        else:
            print "Error: unknown NERSC host '%s'. Using default job submission parameters." % self.launcher.nersc_host
            print "If this does not work, please contact consult@nersc.gov."
            args = super(JobSubmitter_qsub_NERSC, self).SetupPPN(nodes, procs, ppn, use_vis)
        return args

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

class NERSCLauncher(MainLauncher):
    def __init__(self):
        super(NERSCLauncher, self).__init__()
        self.loginnodeport = None
        self.loginnodehost = None
        self.nersc_host = GETENV("NERSC_HOST")
        if not self.nersc_host:
            if self.hostname() == "euclid.nersc.gov":
                self.nersc_host = "euclid"
            elif self.sectorname() == "cvrsvc":
                self.nersc_host = "carver"
            else:
                exit("Cannot determine NERSC host.", 1)

    def Customize(self):
       if self.nersc_host == "edison":
        # ----
        # Edison @ NERSC
        # ----
            # Add GCC libraries to LD_LIBRARY_PATH
            ld_library_path = self.splitpaths(GETENV("LD_LIBRARY_PATH"))
            added_ld_library_paths = ["/opt/gcc/4.8.0/snos/lib64"]
            SETENV("LD_LIBRARY_PATH", self.joinpaths(added_ld_library_paths + ld_library_path))

            # Set Python path
            python_path = self.splitpaths(GETENV("PYTHONPATH"))
            added_python_paths = [self.visitlibdir + "/python/lib/python2.6/lib-dynload"]
            SETENV("PYTHONPATH", self.joinpaths(added_python_paths + python_path))

            # Running the gui on Hopper or a serial engine on a login node are "verboten"
#            if self.sectorname() == "edison":
#                if not self.generalArgs.env and self.generalArgs.exe_name == "gui":
#                    msg = """
#Do not run the VisIt GUI on Hopper. Run it on you local workstation (preferred/best performance) or Euclid!
#
#For more information about running VisIt at NERSC: http://www.nersc.gov/nusers/resources/software/apps/visualization/visit/
#
#"""
#                    exit(msg, 0)
#
#                if self.generalArgs.exe_name == "engine" and not self.parallelArgs.parallel:
#                    exit("The VisIt install on Hopper does not support running a serial engine!", 1)
       elif self.nersc_host == "grace" or self.nersc_host == "hopper":
        # ----
        # Hopper @ NERSC
        # ----
            # Add GCC libraries to LD_LIBRARY_PATH
            ld_library_path = self.splitpaths(GETENV("LD_LIBRARY_PATH"))
            added_ld_library_paths = ["/opt/gcc/4.7.2/snos/lib64"]
            SETENV("LD_LIBRARY_PATH", self.joinpaths(added_ld_library_paths + ld_library_path))

            # Running the gui on Hopper or a serial engine on a login node are "verboten"
            if self.sectorname() == "grace" or self.sectorname() == "hopper":
                if not self.generalArgs.env and self.generalArgs.exe_name == "gui" and not self.visitver[-1] =="b":
                    msg = """
Do not run the VisIt GUI on Hopper. Run it on you local workstation (preferred/best performance) or Euclid!

For more information about running VisIt at NERSC: http://www.nersc.gov/nusers/resources/software/apps/visualization/visit/

"""
                    exit(msg, 0)

                if self.generalArgs.exe_name == "engine" and not self.parallelArgs.parallel:
                    exit("The VisIt install on Hopper does not support running a serial engine!", 1)
            else:
                python_path = self.splitpaths(GETENV("PYTHONPATH"))
                added_python_paths = [self.visitlibdir + "/python/lib/python2.6/lib-dynload"]
                SETENV("PYTHONPATH", self.joinpaths(added_python_paths + python_path))

    #
    # Override the JobSubmitterFactory method so the custom job submitter can
    # be returned.
    #
    def JobSubmitterFactory(self, launch):
        if launch == "aprun":
            return JobSubmitter_aprun_NERSC(self)
        if launch == "mpirun":
            return JobSubmitter_mpirun_NERSC(self)
        elif launch[:4] == "qsub" or launch[:4] == "msub":
            return JobSubmitter_qsub_NERSC(self)
        return super(NERSCLauncher, self).JobSubmitterFactory(launch)

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