"""Analysis backends --- :mod:`MDAnalysis.analysis.backends`
============================================================

.. versionadded:: 2.8.0


The :mod:`backends` module provides :class:`BackendBase` base class to
implement custom execution backends for
:meth:`MDAnalysis.analysis.base.AnalysisBase.run` and its
subclasses.

.. SeeAlso:: :ref:`parallel-analysis`

.. _backends:

Backends
--------

Three built-in backend classes are provided:

* *serial*: :class:`BackendSerial`, that is equivalent to using no
  parallelization and is the default

* *multiprocessing*: :class:`BackendMultiprocessing` that supports
  parallelization via standard Python :mod:`multiprocessing` module
  and uses default :mod:`pickle` serialization

* *dask*: :class:`BackendDask`, that uses the same process-based
  parallelization as :class:`BackendMultiprocessing`, but different
  serialization algorithm via `dask <https://dask.org/>`_ (see `dask
  serialization algorithms
  <https://distributed.dask.org/en/latest/serialization.html>`_ for details)

Classes
-------

"""

import warnings
from typing import Callable
from MDAnalysis.lib.util import is_installed


class BackendBase:
    """Base class for backend implementation.

    Initializes an instance and performs checks for its validity, such as
    ``n_workers`` and possibly other ones.

    Parameters
    ----------
    n_workers : int
        number of workers (usually, processes) over which the work is split

    Examples
    --------
    .. code-block:: python

        from MDAnalysis.analysis.backends import BackendBase

        class ThreadsBackend(BackendBase):
            def apply(self, func, computations):
                from multiprocessing.dummy import Pool

                with Pool(processes=self.n_workers) as pool:
                    results = pool.map(func, computations)
                return results

        import MDAnalysis as mda
        from MDAnalysis.tests.datafiles import PSF, DCD
        from MDAnalysis.analysis.rms import RMSD

        u = mda.Universe(PSF, DCD)
        ref = mda.Universe(PSF, DCD)

        R = RMSD(u, ref)

        n_workers = 2
        backend = ThreadsBackend(n_workers=n_workers)
        R.run(backend=backend, unsupported_backend=True)

    .. warning::
        Using `ThreadsBackend` above will lead to erroneous results, since it
        is an educational example. Do not use it for real analysis.


    .. versionadded:: 2.8.0

    """

    def __init__(self, n_workers: int):
        self.n_workers = n_workers
        self._validate()

    def _get_checks(self):
        """Get dictionary with ``condition: error_message`` pairs that ensure the
        validity of the backend instance

        Returns
        -------
        dict
            dictionary with ``condition: error_message`` pairs that will get
            checked during ``_validate()`` run
        """
        return {
            isinstance(self.n_workers, int)
            and self.n_workers
            > 0: f"n_workers should be positive integer, got {self.n_workers=}",
        }

    def _get_warnings(self):
        """Get dictionary with ``condition: warning_message`` pairs that ensure
        the good usage of the backend instance

        Returns
        -------
        dict
            dictionary with ``condition: warning_message`` pairs that will get
            checked during ``_validate()`` run
        """
        return dict()

    def _validate(self):
        """Check correctness (e.g. ``dask`` is installed if using ``backend='dask'``)
        and good usage (e.g. ``n_workers=1`` if backend is serial) of the backend

        Raises
        ------
        ValueError
            if one of the conditions in :meth:`_get_checks` is ``True``
        """
        for check, msg in self._get_checks().items():
            if not check:
                raise ValueError(msg)
        for check, msg in self._get_warnings().items():
            if not check:
                warnings.warn(msg)

    def apply(self, func: Callable, computations: list) -> list:
        """map function `func` to all tasks in the `computations` list

        Main method that will get called when using an instance of
        ``BackendBase``.  It is equivalent to running ``[func(item) for item in
        computations]`` while using the parallel backend capabilities.

        Parameters
        ----------
        func : Callable
            function to be called on each of the tasks in computations list
        computations : list
            computation tasks to apply function to

        Returns
        -------
        list
            list of results of the function

        """
        raise NotImplementedError


class BackendSerial(BackendBase):
    """A built-in backend that does serial execution of the function, without any
    parallelization.

    Parameters
    ----------
    n_workers : int
        Is ignored in this class, and if ``n_workers`` > 1, a warning will be
        given.


    .. versionadded:: 2.8.0
    """

    def _get_warnings(self):
        """Get dictionary with ``condition: warning_message`` pairs that ensure
        the good usage of the backend instance. Here, it checks if the number
        of workers is not 1, otherwise gives warning.

        Returns
        -------
        dict
            dictionary with ``condition: warning_message`` pairs that will get
            checked during ``_validate()`` run
        """
        return {
            self.n_workers
            == 1: "n_workers is ignored when executing with backend='serial'"
        }

    def apply(self, func: Callable, computations: list) -> list:
        """
        Serially applies `func` to each task object in ``computations``.

        Parameters
        ----------
        func : Callable
            function to be called on each of the tasks in computations list
        computations : list
            computation tasks to apply function to

        Returns
        -------
        list
            list of results of the function
        """
        return [func(task) for task in computations]


class BackendMultiprocessing(BackendBase):
    """A built-in backend that executes a given function using the
    :meth:`multiprocessing.Pool.map <multiprocessing.pool.Pool.map>` method.

    Parameters
    ----------
    n_workers : int
        number of processes in :class:`multiprocessing.Pool
        <multiprocessing.pool.Pool>` to distribute the workload
        between. Must be a positive integer.

    Examples
    --------

    .. code-block:: python

        from MDAnalysis.analysis.backends import BackendMultiprocessing
        import multiprocessing as mp

        backend_obj = BackendMultiprocessing(n_workers=mp.cpu_count())


    .. versionadded:: 2.8.0

    """

    def apply(self, func: Callable, computations: list) -> list:
        """Applies `func` to each object in ``computations`` using `multiprocessing`'s `Pool.map`.

        Parameters
        ----------
        func : Callable
            function to be called on each of the tasks in computations list
        computations : list
            computation tasks to apply function to

        Returns
        -------
        list
            list of results of the function
        """
        from multiprocessing import Pool

        with Pool(processes=self.n_workers) as pool:
            results = pool.map(func, computations)
        return results


class BackendDask(BackendBase):
    """A built-in backend that executes a given function with *dask*.

    Execution is performed with the :func:`dask.compute` function of
    :class:`dask.delayed.Delayed` object (created with
    :func:`dask.delayed.delayed`) with ``scheduler='processes'`` and
    ``chunksize=1`` (this ensures uniform distribution of tasks among
    processes). Requires the `dask package <https://docs.dask.org/en/stable/>`_
    to be `installed <https://docs.dask.org/en/stable/install.html>`_.

    Parameters
    ----------
    n_workers : int
        number of processes in to distribute the workload
        between. Must be a positive integer. Workers are actually
        :class:`multiprocessing.pool.Pool` processes, but they use a different and
        more flexible `serialization protocol
        <https://docs.dask.org/en/stable/phases-of-computation.html#graph-serialization>`_.

    Examples
    --------

    .. code-block:: python

        from MDAnalysis.analysis.backends import BackendDask
        import multiprocessing as mp

        backend_obj = BackendDask(n_workers=mp.cpu_count())


    .. versionadded:: 2.8.0

    """

    def apply(self, func: Callable, computations: list) -> list:
        """Applies `func` to each object in ``computations``.

        Parameters
        ----------
        func : Callable
            function to be called on each of the tasks in computations list
        computations : list
            computation tasks to apply function to

        Returns
        -------
        list
            list of results of the function
        """
        from dask.delayed import delayed
        import dask

        computations = [delayed(func)(task) for task in computations]
        results = dask.compute(
            computations,
            scheduler="processes",
            chunksize=1,
            num_workers=self.n_workers,
        )[0]
        return results

    def _get_checks(self):
        """Get dictionary with ``condition: error_message`` pairs that ensure the
        validity of the backend instance. Here checks if ``dask`` module is
        installed in the environment.

        Returns
        -------
        dict
            dictionary with ``condition: error_message`` pairs that will get
            checked during ``_validate()`` run
        """
        base_checks = super()._get_checks()
        checks = {
            is_installed("dask"): (
                "module 'dask' is missing. Please install 'dask': "
                "https://docs.dask.org/en/stable/install.html"
            )
        }
        return base_checks | checks
