# fmt: off

import collections
from functools import reduce, singledispatch
from typing import (
    Any,
    Dict,
    Iterable,
    List,
    Optional,
    Sequence,
    TypeVar,
    Union,
    overload,
)

import numpy as np
from matplotlib.axes import Axes

from ase.spectrum.dosdata import DOSData, Floats, GridDOSData, Info, RawDOSData
from ase.utils.plotting import SimplePlottingAxes


class DOSCollection(collections.abc.Sequence):
    """Base class for a collection of DOSData objects"""

    def __init__(self, dos_series: Iterable[DOSData]) -> None:
        self._data = list(dos_series)

    def _sample(self,
                energies: Floats,
                width: float = 0.1,
                smearing: str = 'Gauss') -> np.ndarray:
        """Sample the DOS data at chosen points, with broadening

        This samples the underlying DOS data in the same way as the .sample()
        method of those DOSData items, returning a 2-D array with columns
        corresponding to x and rows corresponding to the collected data series.

        Args:
            energies: energy values for sampling
            width: Width of broadening kernel
            smearing: selection of broadening kernel (only "Gauss" is currently
                supported)

        Returns:
            Weights sampled from a broadened DOS at values corresponding to x,
            in rows corresponding to DOSData entries contained in this object
        """

        if len(self) == 0:
            raise IndexError("No data to sample")

        return np.asarray(
            [data._sample(energies, width=width, smearing=smearing)
             for data in self])

    def plot(self,
             npts: int = 1000,
             xmin: float = None,
             xmax: float = None,
             width: float = 0.1,
             smearing: str = 'Gauss',
             ax: Axes = None,
             show: bool = False,
             filename: str = None,
             mplargs: dict = None) -> Axes:
        """Simple plot of collected DOS data, resampled onto a grid

        If the special key 'label' is present in self.info, this will be set
        as the label for the plotted line (unless overruled in mplargs). The
        label is only seen if a legend is added to the plot (i.e. by calling
        `ax.legend()`).

        Args:
            npts, xmin, xmax: output data range, as passed to self.sample_grid
            width: Width of broadening kernel, passed to self.sample_grid()
            smearing: selection of broadening kernel for self.sample_grid()
            ax: existing Matplotlib axes object. If not provided, a new figure
                with one set of axes will be created using Pyplot
            show: show the figure on-screen
            filename: if a path is given, save the figure to this file
            mplargs: additional arguments to pass to matplotlib plot command
                (e.g. {'linewidth': 2} for a thicker line).

        Returns:
            Plotting axes. If "ax" was set, this is the same object.
        """
        return self.sample_grid(npts,
                                xmin=xmin, xmax=xmax,
                                width=width, smearing=smearing
                                ).plot(npts=npts,
                                       xmin=xmin, xmax=xmax,
                                       width=width, smearing=smearing,
                                       ax=ax, show=show, filename=filename,
                                       mplargs=mplargs)

    def sample_grid(self,
                    npts: int,
                    xmin: float = None,
                    xmax: float = None,
                    padding: float = 3,
                    width: float = 0.1,
                    smearing: str = 'Gauss',
                    ) -> 'GridDOSCollection':
        """Sample the DOS data on an evenly-spaced energy grid

        Args:
            npts: Number of sampled points
            xmin: Minimum sampled energy value; if unspecified, a default is
                chosen
            xmax: Maximum sampled energy value; if unspecified, a default is
                chosen
            padding: If xmin/xmax is unspecified, default value will be padded
                by padding * width to avoid cutting off peaks.
            width: Width of broadening kernel, passed to self.sample_grid()
            smearing: selection of broadening kernel, for self.sample_grid()

        Returns:
            (energy values, sampled DOS)
        """
        if len(self) == 0:
            raise IndexError("No data to sample")

        if xmin is None:
            xmin = (min(min(data.get_energies()) for data in self)
                    - (padding * width))
        if xmax is None:
            xmax = (max(max(data.get_energies()) for data in self)
                    + (padding * width))

        return GridDOSCollection(
            [data.sample_grid(npts, xmin=xmin, xmax=xmax, width=width,
                              smearing=smearing)
             for data in self])

    @classmethod
    def from_data(cls,
                  energies: Floats,
                  weights: Sequence[Floats],
                  info: Sequence[Info] = None) -> 'DOSCollection':
        """Create a DOSCollection from data sharing a common set of energies

        This is a convenience method to be used when all the DOS data in the
        collection has a common energy axis. There is no performance advantage
        in using this method for the generic DOSCollection, but for
        GridDOSCollection it is more efficient.

        Args:
            energy: common set of energy values for input data
            weights: array of DOS weights with rows corresponding to different
                datasets
            info: sequence of info dicts corresponding to weights rows.

        Returns:
            Collection of DOS data (in RawDOSData format)
        """

        info = cls._check_weights_and_info(weights, info)

        return cls(RawDOSData(energies, row_weights, row_info)
                   for row_weights, row_info in zip(weights, info))

    @staticmethod
    def _check_weights_and_info(weights: Sequence[Floats],
                                info: Optional[Sequence[Info]],
                                ) -> Sequence[Info]:
        if info is None:
            info = [{} for _ in range(len(weights))]
        else:
            if len(info) != len(weights):
                raise ValueError("Length of info must match number of rows in "
                                 "weights")
        return info

    @overload
    def __getitem__(self, item: int) -> DOSData:
        ...

    @overload  # noqa F811
    def __getitem__(self, item: slice) -> 'DOSCollection':  # noqa F811
        ...

    def __getitem__(self, item):  # noqa F811
        if isinstance(item, int):
            return self._data[item]
        elif isinstance(item, slice):
            return type(self)(self._data[item])
        else:
            raise TypeError("index in DOSCollection must be an integer or "
                            "slice")

    def __len__(self) -> int:
        return len(self._data)

    def _almost_equals(self, other: Any) -> bool:
        """Compare with another DOSCollection for testing purposes"""
        if not isinstance(other, type(self)):
            return False
        elif len(self) != len(other):
            return False
        else:
            return all(a._almost_equals(b) for a, b in zip(self, other))

    def total(self) -> DOSData:
        """Sum all the DOSData in this Collection and label it as 'Total'"""
        data = self.sum_all()
        data.info.update({'label': 'Total'})
        return data

    def sum_all(self) -> DOSData:
        """Sum all the DOSData contained in this Collection"""
        if len(self) == 0:
            raise IndexError("No data to sum")
        elif len(self) == 1:
            data = self[0].copy()
        else:
            data = reduce(lambda x, y: x + y, self)
        return data

    D = TypeVar('D', bound=DOSData)

    @staticmethod
    def _select_to_list(dos_collection: Sequence[D],         # Bug in flakes
                        info_selection: Dict[str, str],      # misses 'D' def
                        negative: bool = False) -> List[D]:  # noqa: F821
        query = set(info_selection.items())

        if negative:
            return [data for data in dos_collection
                    if not query.issubset(set(data.info.items()))]
        else:
            return [data for data in dos_collection
                    if query.issubset(set(data.info.items()))]

    def select(self, **info_selection: str) -> 'DOSCollection':
        """Narrow DOSCollection to items with specified info

        For example, if ::

          dc = DOSCollection([DOSData(x1, y1, info={'a': '1', 'b': '1'}),
                              DOSData(x2, y2, info={'a': '2', 'b': '1'})])

        then ::

          dc.select(b='1')

        will return an identical object to dc, while ::

          dc.select(a='1')

        will return a DOSCollection with only the first item and ::

          dc.select(a='2', b='1')

        will return a DOSCollection with only the second item.

        """

        matches = self._select_to_list(self, info_selection)
        return type(self)(matches)

    def select_not(self, **info_selection: str) -> 'DOSCollection':
        """Narrow DOSCollection to items without specified info

        For example, if ::

          dc = DOSCollection([DOSData(x1, y1, info={'a': '1', 'b': '1'}),
                              DOSData(x2, y2, info={'a': '2', 'b': '1'})])

        then ::

          dc.select_not(b='2')

        will return an identical object to dc, while ::

          dc.select_not(a='2')

        will return a DOSCollection with only the first item and ::

          dc.select_not(a='1', b='1')

        will return a DOSCollection with only the second item.

        """
        matches = self._select_to_list(self, info_selection, negative=True)
        return type(self)(matches)

    # Use typehint *info_keys: str from python3.11+
    def sum_by(self, *info_keys) -> 'DOSCollection':
        """Return a DOSCollection with some data summed by common attributes

        For example, if ::

          dc = DOSCollection([DOSData(x1, y1, info={'a': '1', 'b': '1'}),
                              DOSData(x2, y2, info={'a': '2', 'b': '1'}),
                              DOSData(x3, y3, info={'a': '2', 'b': '2'})])

        then ::

          dc.sum_by('b')

        will return a collection equivalent to ::

          DOSCollection([DOSData(x1, y1, info={'a': '1', 'b': '1'})
                         + DOSData(x2, y2, info={'a': '2', 'b': '1'}),
                         DOSData(x3, y3, info={'a': '2', 'b': '2'})])

        where the resulting contained DOSData have info attributes of
        {'b': '1'} and {'b': '2'} respectively.

        dc.sum_by('a', 'b') on the other hand would return the full three-entry
        collection, as none of the entries have common 'a' *and* 'b' info.

        """

        def _matching_info_tuples(data: DOSData):
            """Get relevent dict entries in tuple form

            e.g. if data.info = {'a': 1, 'b': 2, 'c': 3}
                 and info_keys = ('a', 'c')

                 then return (('a', 1), ('c': 3))
            """
            matched_keys = set(info_keys) & set(data.info)
            return tuple(sorted([(key, data.info[key])
                                 for key in matched_keys]))

        # Sorting inside info matching helps set() to remove redundant matches;
        # combos are then sorted() to ensure consistent output across sessions.
        all_combos = map(_matching_info_tuples, self)
        unique_combos = sorted(set(all_combos))

        # For each key/value combination, perform a select() to obtain all
        # the matching entries and sum them together.
        collection_data = [self.select(**dict(combo)).sum_all()
                           for combo in unique_combos]
        return type(self)(collection_data)

    def __add__(self, other: Union['DOSCollection', DOSData]
                ) -> 'DOSCollection':
        """Join entries between two DOSCollection objects of the same type

        It is also possible to add a single DOSData object without wrapping it
        in a new collection: i.e. ::

          DOSCollection([dosdata1]) + DOSCollection([dosdata2])

        or ::

          DOSCollection([dosdata1]) + dosdata2

        will return ::

          DOSCollection([dosdata1, dosdata2])

        """
        return _add_to_collection(other, self)


@singledispatch
def _add_to_collection(other: Union[DOSData, DOSCollection],
                       collection: DOSCollection) -> DOSCollection:
    if isinstance(other, type(collection)):
        return type(collection)(list(collection) + list(other))
    elif isinstance(other, DOSCollection):
        raise TypeError("Only DOSCollection objects of the same type may "
                        "be joined with '+'.")
    else:
        raise TypeError("DOSCollection may only be joined to DOSData or "
                        "DOSCollection objects with '+'.")


@_add_to_collection.register(DOSData)
def _add_data(other: DOSData, collection: DOSCollection) -> DOSCollection:
    """Return a new DOSCollection with an additional DOSData item"""
    return type(collection)(list(collection) + [other])


class RawDOSCollection(DOSCollection):
    def __init__(self, dos_series: Iterable[RawDOSData]) -> None:
        super().__init__(dos_series)
        for dos_data in self:
            if not isinstance(dos_data, RawDOSData):
                raise TypeError("RawDOSCollection can only store "
                                "RawDOSData objects.")


class GridDOSCollection(DOSCollection):
    def __init__(self, dos_series: Iterable[GridDOSData],
                 energies: Optional[Floats] = None) -> None:
        dos_list = list(dos_series)
        if energies is None:
            if len(dos_list) == 0:
                raise ValueError("Must provide energies to create a "
                                 "GridDOSCollection without any DOS data.")
            self._energies = dos_list[0].get_energies()
        else:
            self._energies = np.asarray(energies)

        self._weights: np.ndarray = np.empty(
            (len(dos_list), len(self._energies)), float,
        )
        self._info = []

        for i, dos_data in enumerate(dos_list):
            if not isinstance(dos_data, GridDOSData):
                raise TypeError("GridDOSCollection can only store "
                                "GridDOSData objects.")
            if (dos_data.get_energies().shape != self._energies.shape
                    or not np.allclose(dos_data.get_energies(),
                                       self._energies)):
                raise ValueError("All GridDOSData objects in GridDOSCollection"
                                 " must have the same energy axis.")
            self._weights[i, :] = dos_data.get_weights()
            self._info.append(dos_data.info)

    def get_energies(self) -> Floats:
        return self._energies.copy()

    def get_all_weights(self) -> Union[Sequence[Floats], np.ndarray]:
        return self._weights.copy()

    def __len__(self) -> int:
        return self._weights.shape[0]

    @overload  # noqa F811
    def __getitem__(self, item: int) -> DOSData:
        ...

    @overload  # noqa F811
    def __getitem__(self, item: slice) -> 'GridDOSCollection':  # noqa F811
        ...

    def __getitem__(self, item):  # noqa F811
        if isinstance(item, int):
            return GridDOSData(self._energies, self._weights[item, :],
                               info=self._info[item])
        elif isinstance(item, slice):
            return type(self)([self[i] for i in range(len(self))[item]])
        else:
            raise TypeError("index in DOSCollection must be an integer or "
                            "slice")

    @classmethod
    def from_data(cls,
                  energies: Floats,
                  weights: Sequence[Floats],
                  info: Sequence[Info] = None) -> 'GridDOSCollection':
        """Create a GridDOSCollection from data with a common set of energies

        This convenience method may also be more efficient as it limits
        redundant copying/checking of the data.

        Args:
            energies: common set of energy values for input data
            weights: array of DOS weights with rows corresponding to different
                datasets
            info: sequence of info dicts corresponding to weights rows.

        Returns:
            Collection of DOS data (in RawDOSData format)
        """

        weights_array = np.asarray(weights, dtype=float)
        if len(weights_array.shape) != 2:
            raise IndexError("Weights must be a 2-D array or nested sequence")
        if weights_array.shape[0] < 1:
            raise IndexError("Weights cannot be empty")
        if weights_array.shape[1] != len(energies):
            raise IndexError("Length of weights rows must equal size of x")

        info = cls._check_weights_and_info(weights, info)

        dos_collection = cls([GridDOSData(energies, weights_array[0])])
        dos_collection._weights = weights_array
        dos_collection._info = list(info)

        return dos_collection

    def select(self, **info_selection: str) -> 'DOSCollection':
        """Narrow GridDOSCollection to items with specified info

        For example, if ::

          dc = GridDOSCollection([GridDOSData(x, y1,
                                              info={'a': '1', 'b': '1'}),
                                  GridDOSData(x, y2,
                                              info={'a': '2', 'b': '1'})])

        then ::

          dc.select(b='1')

        will return an identical object to dc, while ::

          dc.select(a='1')

        will return a DOSCollection with only the first item and ::

          dc.select(a='2', b='1')

        will return a DOSCollection with only the second item.

        """

        matches = self._select_to_list(self, info_selection)
        if len(matches) == 0:
            return type(self)([], energies=self._energies)
        else:
            return type(self)(matches)

    def select_not(self, **info_selection: str) -> 'DOSCollection':
        """Narrow GridDOSCollection to items without specified info

        For example, if ::

          dc = GridDOSCollection([GridDOSData(x, y1,
                                              info={'a': '1', 'b': '1'}),
                                  GridDOSData(x, y2,
                                              info={'a': '2', 'b': '1'})])

        then ::

          dc.select_not(b='2')

        will return an identical object to dc, while ::

          dc.select_not(a='2')

        will return a DOSCollection with only the first item and ::

          dc.select_not(a='1', b='1')

        will return a DOSCollection with only the second item.

        """
        matches = self._select_to_list(self, info_selection, negative=True)
        if len(matches) == 0:
            return type(self)([], energies=self._energies)
        else:
            return type(self)(matches)

    def plot(self,
             npts: int = 0,
             xmin: float = None,
             xmax: float = None,
             width: float = None,
             smearing: str = 'Gauss',
             ax: Axes = None,
             show: bool = False,
             filename: str = None,
             mplargs: dict = None) -> Axes:
        """Simple plot of collected DOS data, resampled onto a grid

        If the special key 'label' is present in self.info, this will be set
        as the label for the plotted line (unless overruled in mplargs). The
        label is only seen if a legend is added to the plot (i.e. by calling
        `ax.legend()`).

        Args:
            npts:
                Number of points in resampled x-axis. If set to zero (default),
                no resampling is performed and the stored data is plotted
                directly.
            xmin, xmax:
                output data range; this limits the resampling range as well as
                the plotting output
            width: Width of broadening kernel, passed to self.sample()
            smearing: selection of broadening kernel, passed to self.sample()
            ax: existing Matplotlib axes object. If not provided, a new figure
                with one set of axes will be created using Pyplot
            show: show the figure on-screen
            filename: if a path is given, save the figure to this file
            mplargs: additional arguments to pass to matplotlib plot command
                (e.g. {'linewidth': 2} for a thicker line).

        Returns:
            Plotting axes. If "ax" was set, this is the same object.
        """

        # Apply defaults if necessary
        npts, width = GridDOSData._interpret_smearing_args(npts, width)

        if npts:
            assert isinstance(width, float)
            dos = self.sample_grid(npts,
                                   xmin=xmin, xmax=xmax,
                                   width=width, smearing=smearing)
        else:
            dos = self

        energies, all_y = dos._energies, dos._weights

        all_labels = [DOSData.label_from_info(data.info) for data in self]

        with SimplePlottingAxes(ax=ax, show=show, filename=filename) as ax:
            self._plot_broadened(ax, energies, all_y, all_labels, mplargs)

        return ax

    @staticmethod
    def _plot_broadened(ax: Axes,
                        energies: Floats,
                        all_y: np.ndarray,
                        all_labels: Sequence[str],
                        mplargs: Optional[Dict]):
        """Plot DOS data with labels to axes

        This is separated into another function so that subclasses can
        manipulate broadening, labels etc in their plot() method."""
        if mplargs is None:
            mplargs = {}

        all_lines = ax.plot(energies, all_y.T, **mplargs)
        for line, label in zip(all_lines, all_labels):
            line.set_label(label)
        ax.legend()

        ax.set_xlim(left=min(energies), right=max(energies))
        ax.set_ylim(bottom=0)