# Copyright (c) 2009-2021 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause
# License.

"""MPI communicator."""

from hoomd import _hoomd
import hoomd

import contextlib


class Communicator(object):
    """MPI communicator.

    Args:
        mpi_comm: Accepts an mpi4py communicator. Use this argument to perform
          many independent hoomd simulations where you communicate between those
          simulations using mpi4py.
        ranks_per_partition (int): (MPI) Number of ranks to include in a
          partition.

    `Communicator` initialize MPI communications for a `hoomd.Simulation`. To
    use MPI, launch your Python script with an MPI launcher (e.g. ``mpirun`` or
    ``mpiexec``). By default, `Communicator` uses all ranks provided by the
    launcher ``num_launch_ranks`` for a single `hoomd.Simulation` object which
    decomposes the state onto that many domains.

    Set ``ranks_per_partition`` to an integer to partition launched ranks into
    ``num_launch_ranks / ranks_per_partition`` communicators, each with their
    own `partition` index. Use this to perform many simulations in parallel, for
    example by using `partition` as an index into an array of state points to
    execute.
    """

    def __init__(self, mpi_comm=None, ranks_per_partition=None):

        # check ranks_per_partition
        if ranks_per_partition is not None:
            if not hoomd.version.mpi_enabled:
                raise RuntimeError(
                    "The ranks_per_partition option is only available in MPI.\n"
                )

        mpi_available = hoomd.version.mpi_enabled

        self.cpp_mpi_conf = None

        # create the specified configuration
        if mpi_comm is None:
            self.cpp_mpi_conf = _hoomd.MPIConfiguration()
        else:
            if not mpi_available:
                raise RuntimeError("mpi_comm is not supported in serial builds")

            handled = False

            # pass in pointer to MPI_Comm object provided by mpi4py
            try:
                import mpi4py
                if isinstance(mpi_comm, mpi4py.MPI.Comm):
                    addr = mpi4py.MPI._addressof(mpi_comm)
                    self.cpp_mpi_conf = \
                        _hoomd.MPIConfiguration._make_mpi_conf_mpi_comm(addr)
                    handled = True
            except ImportError:
                # silently ignore when mpi4py is missing
                pass

            # undocumented case: handle plain integers as pointers to MPI_Comm
            # objects
            if not handled and isinstance(mpi_comm, int):
                self.cpp_mpi_conf = \
                    _hoomd.MPIConfiguration._make_mpi_conf_mpi_comm(mpi_comm)
                handled = True

            if not handled:
                raise RuntimeError(
                    "Invalid mpi_comm object: {}".format(mpi_comm))

        if ranks_per_partition is not None:
            # check validity
            if (self.cpp_mpi_conf.getNRanksGlobal() % ranks_per_partition):
                raise RuntimeError('Total number of ranks is not a multiple of '
                                   'ranks_per_partition.')

            # split the communicator into partitions
            self.cpp_mpi_conf.splitPartitions(ranks_per_partition)

    @property
    def num_ranks(self):
        """int: The number of ranks in this partition.

        When initialized with ``ranks_per_partition=None``, `num_ranks` is equal
        to the ``num_launch_ranks`` set by the MPI launcher. When using
        partitions, `num_ranks` is equal to ``ranks_per_partition``.

        Note:
            Returns 1 in builds with ENABLE_MPI=off.
        """
        if hoomd.version.mpi_enabled:
            return self.cpp_mpi_conf.getNRanks()
        else:
            return 1

    @property
    def rank(self):
        """int: The current rank within the partition.

        Note:
            Returns 0 in builds with ENABLE_MPI=off.
        """
        if hoomd.version.mpi_enabled:
            return self.cpp_mpi_conf.getRank()
        else:
            return 0

    @property
    def num_partitions(self):
        """int: The number of partitions in this execution.

        Create partitions with the ``ranks_per_partition`` argument on
        initialization. Then, the number of partitions is
        ``num_launch_ranks / ranks_per_partition``.

        Note:
            Returns 1 in builds with ENABLE_MPI=off.
        """
        if hoomd.version.mpi_enabled:
            return self.cpp_mpi_conf.getNPartitions()
        else:
            return 1

    @property
    def partition(self):
        """int: The current partition.

        Note:
            Returns 0 in builds with ENABLE_MPI=off.
        """
        if hoomd.version.mpi_enabled:
            return self.cpp_mpi_conf.getPartition()
        else:
            return 0

    def barrier_all(self):
        """Perform a MPI barrier synchronization across all ranks.

        Note:
            Does nothing in builds with ENABLE_MPI=off.
        """
        if hoomd.version.mpi_enabled:
            _hoomd.mpi_barrier_world()

    def barrier(self):
        """Perform a barrier synchronization across all ranks in the partition.

        Note:
            Does nothing in builds with ENABLE_MPI=off.
        """
        if hoomd.version.mpi_enabled:
            self.cpp_mpi_conf.barrier()

    @contextlib.contextmanager
    def localize_abort(self):
        """Localize MPI_Abort to this partition.

        HOOMD calls ``MPI_Abort`` to tear down all running MPI processes
        whenever there is an uncaught exception. By default, this will abort the
        entire MPI execution. When using partitions, an uncaught exception on
        one partition will therefore abort all of them.

        Use the return value of :py:meth:`localize_abort()` as a context manager
        to tell HOOMD that all operations within the context will use only
        that MPI communicator so that an uncaught exception in one partition
        will only abort that partition and leave the others running.
        """
        global _current_communicator
        prev = _current_communicator

        _current_communicator = self
        yield None
        _current_communicator = prev


# store the "current" communicator to be used for MPI_Abort calls. This defaults
# to the world communicator, but users can opt in to a more specific
# communicator using the Device.localize_abort context manager
_current_communicator = Communicator()

[((51, 32, 51, 57), 'hoomd._hoomd.MPIConfiguration', '_hoomd.MPIConfiguration', ({}, {}), '()', False, 'from hoomd import _hoomd\n'), ((153, 12, 153, 38), 'hoomd._hoomd.mpi_barrier_world', '_hoomd.mpi_barrier_world', ({}, {}), '()', False, 'from hoomd import _hoomd\n'), ((74, 20, 74, 77), 'hoomd._hoomd.MPIConfiguration._make_mpi_conf_mpi_comm', '_hoomd.MPIConfiguration._make_mpi_conf_mpi_comm', ({(74, 68, 74, 76): 'mpi_comm'}, {}), '(mpi_comm)', False, 'from hoomd import _hoomd\n'), ((62, 27, 62, 58), 'mpi4py.MPI._addressof', 'mpi4py.MPI._addressof', ({(62, 49, 62, 57): 'mpi_comm'}, {}), '(mpi_comm)', False, 'import mpi4py\n'), ((64, 24, 64, 77), 'hoomd._hoomd.MPIConfiguration._make_mpi_conf_mpi_comm', '_hoomd.MPIConfiguration._make_mpi_conf_mpi_comm', ({(64, 72, 64, 76): 'addr'}, {}), '(addr)', False, 'from hoomd import _hoomd\n')]