#!/usr/bin/env python
##@namespace produtil.mpi_impl.mpi_impl_base
# Utilities like CMDFGen to simplify adding new MPI implementations to the 
# produtil.run suite of modules.
#
# This module contains classes and functions to assist developers in
# extending the functionality of the produtil.mpi_impl package.  The
# main highlight is the CMDFGen, which generates command files.  Some
# MPI implementations, and the mpiserial program, want to read a file
# with one line per MPI rank telling what program to run on each rank.
# For example, LSF+IBMPE and LoadLeveler+IBMPE work this way if one
# wants to run different programs on different ranks.

import tempfile,stat,os, logging

module_logger=logging.getLogger('produtil.mpi_impl')

class MPIConfigError(Exception): 
    """!Base class of MPI configuration exceptions."""
class WrongMPI(MPIConfigError): 
    """!Unused: raised when the wrong MPI implementation is accessed.  """
class MPISerialMissing(MPIConfigError):
    """!Raised when the mpiserial program is required, but is missing."""
class MPIAllRanksError(MPIConfigError):
    """!Raised when the allranks=True keyword is sent to mpirun or mpirunner,
but the MPI program specification has more than one rank."""
class MPIMixed(MPIConfigError):
    """!Thrown to indicate serial and parallel processes are being mixed in a single mpi_comm_world."""
class MPIDisabled(MPIConfigError):
    """!Thrown to MPI is not supported."""
class OpenMPDisabled(MPIConfigError):
    """!Raised when OpenMP is not supported by the present implementation."""
class CMDFGen(object):
    """!Generates files with one line per MPI rank, telling what
    program to run on each rank.

    This class is used to generate command files for mpiserial, poe or
    mpirun.lsf.  Command files are files with one MPI rank per line
    containing a shell command to run for that rank.  Generally the
    input (lines) is generated by the to_arglist function in a
    subclass of produtil.mpiprog.MPIRanksBase.  See the
    produtil.mpi_impl.mpirun_lsf for an example of how to use this."""
    def __init__(self,base,lines,cmd_envar='SCR_CMDFILE',
                 model_envar=None,filename_arg=False,**kwargs):
        """!CMDFGen constructor
        
        @param base type of command file being generated.  See below.
        @param lines the command file contents as a list of strings, one per line
        @param cmd_envar environment variable to set to the command file path
        @param model_envar environment variable to set to "MPMD" 
        @param kwargs Sets the command file name.  See below.
        @param filename_arg If True, the name of the command file is appended to the program argument list.

        The command file is generated from
        tempfile.NamedTemporaryFile, passing several arguments from
        kwargs, if provided, or suitable defaults otherwise.  There
        are several arguments used.  In all cases, replace "base" with
        the contents of the @c base argument:

        * base_suffix --- temporary file suffix (default: "base.")
        * base_prefix --- temporary file prefix (default: ".cmdf")
        * base_tempdir --- directory in which to create the file

        @bug The base_suffix keyword is used for both the suffix and prefix"""
        assert(base is not None)
        assert(isinstance(lines,list))
        assert(len(lines)>0)
        assert(isinstance(lines[0],str))
        assert(len(lines[0])>0)
        self.filename=kwargs.get(str(base),None)
        self.tmpprefix=kwargs.get('%s_suffix'%(base,),'%s.'%(base,))
        self.tmpsuffix=kwargs.get('%s_suffix'%(base,),'.cmdf')
        self.tmpdir=kwargs.get('%s_tmpdir'%(base,),'.')
        self.cmd_envar=cmd_envar
        self.model_envar=model_envar
        self.filename_arg=filename_arg
        out='\n'.join(lines)
        if len(out)>0:
            out+='\n'
        self.cmdf_contents=out
        return
    ##@var filename
    # command file's filename

    ##@var tmpprefix 
    # temporary file prefix
    
    ##@var tmpsuffix
    # temporary file suffix

    ##@var tmpdir
    # temporary file directory

    ##@var cmd_envar
    # Environment variable to set telling the path to the
    # command file

    ##@var model_envar
    # Environment variable to set to "MPMD"

    ##@var cmdf_contents
    # String containing the command file contents.

    def _add_more_vars(self,envars,logger):
        """!Adds additional environment variables to the envars dict,
        needed to configure the MPI implementation correctly.  This is
        used to set MP_PGMMODEL="MPMD" if the constructor receives
        model_envar="MP_PGMMODEL".

        @param envars[out] the dict to modify
        @param logger a logging.Logger for log messages"""
        if self.model_envar is not None:
            if logger is not None:
                logger.info('Set %s="MPMD"'%(self.model_envar,))
            envars[self.model_envar]='MPMD'
    def __call__(self,runner,logger=None):
        """!Adds the environment variables to @c runner and creates the command file.

        @param[out] runner A produtil.prog.Runner to modify
        @param logger a logging.Logger for log messages"""
        if logger is None: logger=module_logger
        if self.filename is not None:
            with open(self.filename,'wt') as f:
                f.write(self.cmdf_contents)
            if logger is not None:
                logger.info('Write command file to %s'%(repr(filename),))
            kw={self.cmd_envar: self.filename}
            self._add_more_vars(kw,logger)
            if logger is not None:
                for k,v in kw.items():
                    logger.info('Set %s=%s'%(k,repr(v)))
            if self.filename_arg:
                runner=runner[self.filename]
            return runner.env(**kw)
        else:
            with tempfile.NamedTemporaryFile(mode='wt',suffix=self.tmpsuffix,
                    prefix=self.tmpprefix,dir=self.tmpdir,delete=False) as t:
                if logger is not None:
                    logger.info('Write command file to %s'%(repr(t.name),))
                t.write(self.cmdf_contents)
                # Make the file read-only and readable for everyone:
                os.fchmod(t.fileno(),stat.S_IRUSR|stat.S_IRGRP|stat.S_IROTH)
                kw={self.cmd_envar: t.name}
                self._add_more_vars(kw,logger)
                if logger is not None:
                    for k,v in kw.items():
                        logger.info('Set %s=%s'%(k,repr(v)))
                runner.env(**kw)
                if self.filename_arg:
                    runner=runner[t.name]
            return runner