"""Analysis results and their aggregation --- :mod:`MDAnalysis.analysis.results`
================================================================================

Module introduces two classes, :class:`Results` and :class:`ResultsGroup`,
used for storing and aggregating data in
:meth:`MDAnalysis.analysis.base.AnalysisBase.run()`, respectively.


Classes
-------

The :class:`Results` class is an extension of a built-in dictionary
type, that holds all assigned attributes in :attr:`self.data` and 
allows for access either via dict-like syntax, or via class-like syntax:

.. code-block:: python

    from MDAnalysis.analysis.results import Results
    r = Results()
    r.array = [1, 2, 3, 4]
    assert r['array'] == r.array == [1, 2, 3, 4]


The :class:`ResultsGroup` can merge multiple :class:`Results` objects.
It is mainly used by :class:`MDAnalysis.analysis.base.AnalysisBase` class, 
that uses :meth:`ResultsGroup.merge()` method to aggregate results from
multiple workers, initialized during a parallel run:

.. code-block:: python

    from MDAnalysis.analysis.results import Results, ResultsGroup
    import numpy as np
    
    r1, r2 = Results(), Results()
    r1.masses = [1, 2, 3, 4, 5]
    r2.masses = [0, 0, 0, 0]
    r1.vectors = np.arange(10).reshape(5, 2)
    r2.vectors = np.arange(8).reshape(4, 2)

    group = ResultsGroup(
        lookup = {
            'masses': ResultsGroup.flatten_sequence,
            'vectors': ResultsGroup.ndarray_vstack
            }
        )

    r = group.merge([r1, r2])
    assert r.masses == list((*r1.masses, *r2.masses))
    assert (r.vectors == np.vstack([r1.vectors, r2.vectors])).all()
"""

from collections import UserDict
import numpy as np
from typing import Callable, Sequence


class Results(UserDict):
    r"""Container object for storing results.

    :class:`Results` are dictionaries that provide two ways by which values
    can be accessed: by dictionary key ``results["value_key"]`` or by object
    attribute, ``results.value_key``. :class:`Results` stores all results
    obtained from an analysis after calling :meth:`~AnalysisBase.run()`.

    The implementation is similar to the :class:`sklearn.utils.Bunch`
    class in `scikit-learn`_.

    .. _`scikit-learn`: https://scikit-learn.org/
    .. _`sklearn.utils.Bunch`: https://scikit-learn.org/stable/modules/generated/sklearn.utils.Bunch.html

    Raises
    ------
    AttributeError
        If an assigned attribute has the same name as a default attribute.

    ValueError
        If a key is not of type ``str`` and therefore is not able to be
        accessed by attribute.

    Examples
    --------
    >>> from MDAnalysis.analysis.base import Results
    >>> results = Results(a=1, b=2)
    >>> results['b']
    2
    >>> results.b
    2
    >>> results.a = 3
    >>> results['a']
    3
    >>> results.c = [1, 2, 3, 4]
    >>> results['c']
    [1, 2, 3, 4]


    .. versionadded:: 2.0.0

    .. versionchanged:: 2.8.0
        Moved :class:`Results` to :mod:`MDAnalysis.analysis.results`
    """

    def _validate_key(self, key):
        if key in dir(self):
            raise AttributeError(
                f"'{key}' is a protected dictionary attribute"
            )
        elif isinstance(key, str) and not key.isidentifier():
            raise ValueError(f"'{key}' is not a valid attribute")

    def __init__(self, *args, **kwargs):
        kwargs = dict(*args, **kwargs)
        if "data" in kwargs.keys():
            raise AttributeError(f"'data' is a protected dictionary attribute")
        self.__dict__["data"] = {}
        self.update(kwargs)

    def __setitem__(self, key, item):
        self._validate_key(key)
        super().__setitem__(key, item)

    def __setattr__(self, attr, val):
        if attr == "data":
            super().__setattr__(attr, val)
        else:
            self.__setitem__(attr, val)

    def __getattr__(self, attr):
        try:
            return self[attr]
        except KeyError as err:
            raise AttributeError(
                f"'Results' object has no attribute '{attr}'"
            ) from err

    def __delattr__(self, attr):
        try:
            del self[attr]
        except KeyError as err:
            raise AttributeError(
                f"'Results' object has no attribute '{attr}'"
            ) from err

    def __getstate__(self):
        return self.data

    def __setstate__(self, state):
        self.data = state


class ResultsGroup:
    """
    Management and aggregation of results stored in :class:`Results` instances.

    A :class:`ResultsGroup` is an optional description for :class:`Result` "dictionaries"
    that are used in analysis classes based on :class:`AnalysisBase`. For each *key* in a
    :class:`Result` it describes how multiple pieces of the data held under the key are
    to be aggregated. This approach is necessary when parts of a trajectory are analyzed
    independently (e.g., in parallel) and then need to me merged (with :meth:`merge`) to
    obtain a complete data set.

    Parameters
    ----------
    lookup : dict[str, Callable], optional
        aggregation functions lookup dict, by default None

    Examples
    --------

    .. code-block:: python

        from MDAnalysis.analysis.results import ResultsGroup, Results
        group = ResultsGroup(lookup={'mass': ResultsGroup.float_mean})
        obj1 = Results(mass=1)
        obj2 = Results(mass=3)
        assert {'mass': 2.0} == group.merge([obj1, obj2])


    .. code-block:: python

        # you can also set `None` for those attributes that you want to skip
        lookup = {'mass': ResultsGroup.float_mean, 'trajectory': None}
        group = ResultsGroup(lookup)
        objects = [Results(mass=1, skip=None), Results(mass=3, skip=object)]
        assert group.merge(objects, require_all_aggregators=False) == {'mass': 2.0}

    .. versionadded:: 2.8.0
    """

    def __init__(self, lookup: dict[str, Callable] = None):
        self._lookup = lookup

    def merge(
        self, objects: Sequence[Results], require_all_aggregators: bool = True
    ) -> Results:
        """Merge multiple Results into a single Results instance.

        Merge multiple :class:`Results` instances into a single one, using the
        `lookup` dictionary to determine the appropriate aggregator functions for
        each named results attribute. If the resulting object only contains a single
        element, it just returns it without using any aggregators.

        Parameters
        ----------
        objects : Sequence[Results]
            Multiple :class:`Results` instances with the same data attributes.
        require_all_aggregators : bool, optional
            if True, raise an exception when no aggregation function for a
            particular argument is found. Allows to skip aggregation for the
            parameters that aren't needed in the final object --
            see :class:`ResultsGroup`.

        Returns
        -------
        Results
            merged :class:`Results`

        Raises
        ------
        ValueError
            if no aggregation function for a key is found and ``require_all_aggregators=True``
        """
        if len(objects) == 1:
            merged_results = objects[0]
            return merged_results

        merged_results = Results()
        for key in objects[0].keys():
            agg_function = self._lookup.get(key, None)
            if agg_function is not None:
                results_of_t = [obj[key] for obj in objects]
                merged_results[key] = agg_function(results_of_t)
            elif require_all_aggregators:
                raise ValueError(f"No aggregation function for {key=}")
        return merged_results

    @staticmethod
    def flatten_sequence(arrs: list[list]):
        """Flatten a list of lists into a list

        Parameters
        ----------
        arrs : list[list]
            list of lists

        Returns
        -------
        list
            flattened list
        """
        return [item for sublist in arrs for item in sublist]

    @staticmethod
    def ndarray_sum(arrs: list[np.ndarray]):
        """sums an ndarray along ``axis=0``

        Parameters
        ----------
        arrs : list[np.ndarray]
            list of input arrays. Must have the same shape.

        Returns
        -------
        np.ndarray
            sum of input arrays
        """
        return np.array(arrs).sum(axis=0)

    @staticmethod
    def ndarray_mean(arrs: list[np.ndarray]):
        """calculates mean of input ndarrays along ``axis=0``

        Parameters
        ----------
        arrs : list[np.ndarray]
            list of input arrays. Must have the same shape.

        Returns
        -------
        np.ndarray
            mean of input arrays
        """
        return np.array(arrs).mean(axis=0)

    @staticmethod
    def float_mean(floats: list[float]):
        """calculates mean of input float values

        Parameters
        ----------
        floats : list[float]
            list of float values

        Returns
        -------
        float
            mean value
        """
        return np.array(floats).mean()

    @staticmethod
    def ndarray_hstack(arrs: list[np.ndarray]):
        """Performs horizontal stack of input arrays

        Parameters
        ----------
        arrs : list[np.ndarray]
            input numpy arrays

        Returns
        -------
        np.ndarray
            result of stacking
        """
        return np.hstack(arrs)

    @staticmethod
    def ndarray_vstack(arrs: list[np.ndarray]):
        """Performs vertical stack of input arrays

        Parameters
        ----------
        arrs : list[np.ndarray]
            input numpy arrays

        Returns
        -------
        np.ndarray
            result of stacking
        """
        return np.vstack(arrs)