import math

###############################################################################
# Class: JobSubmitter_sbatch_NERSC
#
# Purpose:    Custom "sbatch" job submitter for NERSC.
#
# Programmer: David Camp
# Date:       Thu Oct  1 08:39:34 PDT 2015
#
# Modifications:
#   Brad Whitlock, Wed Sep 14 15:44:46 PDT 2016
#   Implement socket relay hack as for qsub.
#
###############################################################################

class JobSubmitter_sbatch_NERSC(JobSubmitter_sbatch):
    def __init__(self, launcher):
        super(JobSubmitter_sbatch_NERSC, self).__init__(launcher)

    def CreateFilename(self):
        tdate = time.asctime()[11:19]
        tuser = self.launcher.username()
        return os.path.join(GETENV("HOME"), "visit.%s.%s" % (tuser, tdate))

###############################################################################
# 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
    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", "/opt/cray/alps/5.2.1-2.0502.9072.13.1.gem/bin/aprun"]

###############################################################################
# Class: JobSubmitter_qsub_NERSC
#
# Purpose:    Custom "qsub" job submitter for NERSC.
#
# Programmer: Brad Whitlock
# Date:       Thu May 17 14:22:04 PDT 2012
#
# Modifications:
#   Eric Brugger, Wed Oct 22 16:26:54 PDT 2014
#   Modify the launcher so that -ppn gets passed to mpiexec and mpirun.
#
###############################################################################

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):
        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 == "edison":
            return ["/opt/torque/default/bin/qsub"]
        elif self.launcher.nersc_host == "hopper":
            return ["/usr/common/usg/bin/qsub"]
        elif self.launcher.nersc_host == "grace":
            return ["/usr/common/nsg/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.
        if self.launcher.nersc_host == "edison":
            return ["env", "CRAY_ROOTFS=DSL", "/opt/cray/alps/default/bin/aprun"]
        else:
            return ["env", "CRAY_ROOTFS=DSL", "/opt/cray/alps/5.2.1-2.0502.9072.13.1.gem/bin/aprun"]

    def TFileSetup(self, tfile):
        tfile.write("cd %s\n" % os.path.abspath(os.curdir))
        tfile.write("ulimit -c 0\n")
        # Compute nodes could not talk with login node.
        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:
                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]
        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:
#   Eric Brugger, Wed Oct 22 16:26:54 PDT 2014
#   Modify the launcher so that -ppn gets passed to mpiexec and mpirun.
#
#   David Camp, Thu Oct  1 08:39:34 PDT 2015
#   Added sbatch for Cori system.
#
###############################################################################

class NERSCLauncher(MainLauncher):
    def __init__(self):
        super(NERSCLauncher, self).__init__()
        self.loginnodeport = None
        self.loginnodehost = None
        self.hosttoip = True
        self.nersc_host = GETENV("NERSC_HOST")
        if not self.nersc_host:
            exit("Cannot determine NERSC host.", 1)

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

            # Running the gui on Cori or a serial engine on a login node are "verboten"
            if self.sectorname() == "cori":
                if not self.generalArgs.env and self.generalArgs.exe_name == "gui":
                    msg = """
Do not run the VisIt GUI on Cori. Run it on your local workstation (preferred/best performance)!

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 Edison does not support running a serial engine on the login nodes!", 1)
       elif 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.9.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 Edison 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 Edison. Run it on your local workstation (preferred/best performance)!

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 Edison does not support running a serial engine on the login nodes!", 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.8.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 your local workstation (preferred/best performance)!

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 on the login nodes!", 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[:6] == "sbatch":
            return JobSubmitter_sbatch_NERSC(self)
        elif launch[:4] == "qsub" or launch[:4] == "msub":
            return JobSubmitter_qsub_NERSC(self)
        elif launch == "aprun":
            return JobSubmitter_aprun_NERSC(self)
        return super(NERSCLauncher, self).JobSubmitterFactory(launch)

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