"""Implements the :py:class:`.ColorManager` which simplifies working with
:py:class:`matplotlib.colors.Colormap` and related objects."""

import copy
import logging
from typing import Union

import matplotlib as mpl
import matplotlib.colors
import numpy as np
from matplotlib.colors import to_rgb

log = logging.getLogger(__name__)

# -----------------------------------------------------------------------------

NORMS = {
    "Normalize": mpl.colors.Normalize,
    "BoundaryNorm": mpl.colors.BoundaryNorm,
    "CenteredNorm": mpl.colors.CenteredNorm,
    "NoNorm": mpl.colors.NoNorm,
    "LogNorm": mpl.colors.LogNorm,
    "PowerNorm": mpl.colors.PowerNorm,
    "SymLogNorm": mpl.colors.SymLogNorm,
    "TwoSlopeNorm": mpl.colors.TwoSlopeNorm,
    "FuncNorm": mpl.colors.FuncNorm,
}
"""matplotlib color normalizations supported by the :py:class:`.ColorManager`.
See the :py:mod:`matplotlib.colors` module for more information.
"""


# -----------------------------------------------------------------------------


class ColorManager:
    """Custom color manager which provides an interface to the
    :py:mod:`matplotlib.colors` module and aims to simplify working with
    colormaps, colorbars, and different normalizations.
    """

    _NORMS_NOT_SUPPORTING_VMIN_VMAX = ("BoundaryNorm", "CenteredNorm")
    """Names of norms that do not support getting passed ``vmin`` & ``vmax``"""

    def __init__(
        self,
        *,
        cmap: Union[str, dict, mpl.colors.Colormap] = None,
        norm: Union[str, dict, mpl.colors.Normalize] = None,
        labels: dict = None,
        vmin: float = None,
        vmax: float = None,
    ):
        """Initializes a :py:class:`.ColorManager` by building the colormap,
        the norm, and the colorbar labels.

        Examples (using a YAML representation):

        .. literalinclude:: ../../tests/cfg/color_manager_cfg.yml
            :language: yaml
            :start-after: ### START -- labels
            :end-before:  ### END ---- labels
            :dedent: 2

        A segmented colormap with custom bin edges:

        .. literalinclude:: ../../tests/cfg/color_manager_cfg.yml
            :language: yaml
            :start-after: ### START -- intervals
            :end-before:  ### END ---- intervals
            :dedent: 2

        Specifying color and alpha value:

        .. literalinclude:: ../../tests/cfg/color_manager_cfg.yml
            :language: yaml
            :start-after: ### START -- advanced_syntax
            :end-before:  ### END ---- advanced_syntax
            :dedent: 2

        A continous colormap with linearly interpolated colors at specific
        positions:

        .. literalinclude:: ../../tests/cfg/color_manager_cfg.yml
            :language: yaml
            :start-after: ### START -- continuous_with_boundaries
            :end-before:  ### END ---- continuous_with_boundaries
            :dedent: 2

        Furthermore, YAML *tags* can be used to generate colormaps or norms in
        places where the color manager is not integrated but a corresponding
        :py:class:`matplotlib.colors.Colormap` object or norm object is
        accepted / required:

        .. literalinclude:: ../../tests/cfg/color_manager_cfg.yml
            :language: yaml
            :start-after: ### START -- yaml_constructors
            :end-before:  ### END ---- yaml_constructors
            :dedent: 2

        Args:
            cmap (Union[str, dict, matplotlib.colors.Colormap], optional):
                The colormap. If it is a string, it must name a registered
                colormap. If it is a dict, the following arguments are
                available:

                name (str, optional):
                    Name of a registered colormap.
                from_values (Union[dict, list], optional):
                    Dict of colors keyed by bin-specifier. If given, ``name``
                    is ignored and a discrete colormap is created from the list
                    of specified colors. The ``norm`` is then set to
                    :py:class:`matplotlib.colors.BoundaryNorm`.

                    The bins can be specified either by bin-centers (Scalar) or
                    by bin-intervals (2-tuples). For the former, the deduced
                    bin-edges are assumed halfway between the bin-centers. For
                    the latter, the given intervals must be pairwise connected.
                    In both cases, the bins must monotonically increase.

                    If a list of colors is passed they are automatically
                    assigned to the bin-centers ``[0, 1, 2, ...]``.

                    Alternatively, a continuous, linearly interpolated colormap
                    can be generated by setting the ``continuous`` flag, see
                    below. This will construct a
                    :py:class:`matplotlib.colors.LinearSegmentedColormap`.
                    In such a case, keys in ``from_values`` can only be scalar,
                    bin *intervals* cannot be specified.
                continuous (bool, optional): If True, will interpret the
                    ``from_values`` data as specifying points between which a
                    linear interpolation is carried out. Will create a
                    :py:class:`matplotlib.colors.LinearSegmentedColormap`.
                under (Union[str, dict], optional): Passed on to
                    :py:meth:`matplotlib.colors.Colormap.set_under`
                over (Union[str, dict], optional): Passed on to
                    :py:meth:`matplotlib.colors.Colormap.set_over`
                bad (Union[str, dict], optional): Passed on to
                    :py:meth:`matplotlib.colors.Colormap.set_bad`
                placeholder_color (optional): ``None`` values in
                    ``from_values`` are replaced with this color
                    (default: white).
                further arguments (optional): passed on to colormap
                    initialization.

            norm (Union[str, dict, matplotlib.colors.Normalize], optional):
                The norm that is applied for the color-mapping. If it is a
                string, the matching norm in :py:mod:`matplotlib.colors`
                is created with default values.
                If it is a dict, the ``name`` entry specifies the norm and all
                further entries are passed to its constructor.
                Overwritten if a discrete colormap is specified via
                ``cmap.from_values``.
            labels (Union[dict, list], optional): Colorbar tick-labels keyed by
                tick position. If a list of labels is passed they are
                automatically assigned to the positions [0, 1, 2, ...].
            vmin (float, optional): The lower bound of the color-mapping.
                Ignored if norm is *BoundaryNorm*.
            vmax (float, optional): The upper bound of the color-mapping.
                Ignored if norm is *BoundaryNorm*.
        """
        cmap_kwargs = {}
        norm_kwargs = {}

        # If Colormap instance is given, set it directly. Otherwise parse the
        # cmap_kwargs below.
        if isinstance(cmap, mpl.colors.Colormap):
            self._cmap = cmap

        elif isinstance(cmap, str) or cmap is None:
            cmap_kwargs = dict(name=cmap)

        else:
            cmap_kwargs = copy.deepcopy(cmap)

        # If Normalize instance is given, set it directly. Otherwise parse the
        # norm_kwargs below.
        if isinstance(norm, mpl.colors.Normalize):
            self._norm = norm

        elif isinstance(norm, str) or norm is None:
            norm_kwargs = dict(name=norm)

        else:
            norm_kwargs = copy.deepcopy(norm)

        # Parse configuration for custom color mapping from values, which can
        # be either discrete or continuous (interpolated between colors)
        if "from_values" in cmap_kwargs:
            mapping = cmap_kwargs.pop("from_values")

            # Get the placeholder color. If not given, set white as default.
            _placeholder_color = cmap_kwargs.pop("placeholder_color", "w")

            if isinstance(mapping, list):
                mapping = {k: v for k, v in enumerate(mapping)}

            mapping = {
                k: (v if v is not None else _placeholder_color)
                for k, v in mapping.items()
            }

            # Distinguish between continous and discrete case
            if cmap_kwargs.pop("continuous", False):
                # Get the colordict used to generate the continuous colormap,
                # complying to the interface of LinearSegmentedColormap.
                # Replace all None entries by the placeholder color.
                cdict = dict()
                for num, col in enumerate(("red", "green", "blue")):
                    cdict[col] = [
                        [
                            k,
                            to_rgb(v)[num],
                            to_rgb(v)[num],
                        ]
                        for k, v in mapping.items()
                    ]
                mapping = cdict
                cmap_kwargs["name"] = "LinearSegmentedColormap"
                cmap_kwargs["segmentdata"] = mapping

                log.remark("Configuring a linear colormap 'from values'. ")

            else:
                # Discrete, binned case
                cmap_kwargs["name"] = "ListedColormap"
                cmap_kwargs["colors"] = list(mapping.values())
                norm_kwargs = {
                    "name": "BoundaryNorm",
                    "ncolors": len(mapping),
                    "boundaries": self._parse_boundaries(list(mapping.keys())),
                }
                log.remark(
                    "Configuring a discrete colormap 'from values'. "
                    "Overwriting 'norm' to BoundaryNorm with %s colors.",
                    norm_kwargs["ncolors"],
                )

        # BoundaryNorm has no vmin/vmax argument
        if (
            norm_kwargs.get("name", None)
            not in self._NORMS_NOT_SUPPORTING_VMIN_VMAX
        ):
            norm_kwargs["vmin"] = vmin
            norm_kwargs["vmax"] = vmax

            log.remark("norm.vmin and norm.vmax set to %s and %s", vmin, vmax)

        # Parse shortcut notation
        if isinstance(labels, list):
            labels = {k: v for k, v in enumerate(labels)}

        self._labels = copy.deepcopy(labels)

        # Set cmap and norm if not done already
        if cmap_kwargs:
            self._cmap = self._create_cmap(**cmap_kwargs)

        if norm_kwargs:
            self._norm = self._create_norm(**norm_kwargs)

    @property
    def cmap(self) -> mpl.colors.Colormap:
        """Returns the constructed colormap object"""
        return self._cmap

    @property
    def norm(self) -> mpl.colors.Normalize:
        """Returns the constructed normalization object"""
        return self._norm

    @property
    def labels(self) -> Union[dict, list]:
        """A dict or list of colorbar labels"""
        return self._labels

    # .........................................................................

    def map_to_color(self, X: Union[float, np.ndarray]):
        """Maps the input data to color(s) by applying both norm and colormap.

        Args:
            X (Union[float, numpy.ndarray]): Data value(s) to convert to RGBA.

        Returns:
            Tuple of RGBA values if X is scalar, otherwise an array of RGBA
            values with a shape of ``X.shape + (4, )``.
        """
        return self.cmap(self.norm(X))

    def create_cbar(
        self,
        mappable: "matplotlib.cm.ScalarMappable",
        *,
        fig: "matplotlib.figure.Figure" = None,
        ax: "matplotlib.axes.Axes" = None,
        label: str = None,
        label_kwargs: dict = None,
        tick_params: dict = None,
        **cbar_kwargs,
    ) -> "matplotlib.colorbar.Colorbar":
        """Creates a colorbar of a given mappable

        Args:
            mappable (matplotlib.cm.ScalarMappable): The mappable that is to be
                described by the colorbar.
            fig (matplotlib.figure.Figure, optional): The figure; if not
                given, will use the current figure as determined by
                :py:func:`~matplotlib.pyplot.gcf`.
            ax (matplotlib.axes.Axes, optional): The axes; if not given, will
                use the one given by :py:meth:`matplotlib.figure.Figure.gca`.
            label (str, optional): A label for the colorbar
            label_kwargs (dict, optional): Additional parameters passed to
                :py:meth:`matplotlib.colorbar.Colorbar.set_label`
            tick_params (dict, optional): Set colorbar tick parameters via the
                :py:meth:`matplotlib.axes.Axes.tick_params` method of the
                :py:class:`matplotlib.colorbar.Colorbar` axes.
            **cbar_kwargs: Passed on to
                :py:meth:`matplotlib.figure.Figure.colorbar`

        Returns:
            matplotlib.colorbar.Colorbar: The created colorbar object
        """
        import matplotlib.pyplot as plt

        if fig is None:
            fig = plt.gcf()

        if ax is None:
            ax = fig.gca()

        cb = fig.colorbar(mappable, ax=ax, **cbar_kwargs)

        # Colorbar label
        if label:
            cb.set_label(label=label, **(label_kwargs if label_kwargs else {}))

        # Ticks and tick labels
        if self.labels is not None:
            cb.set_ticks(list(self.labels.keys()))
            cb.set_ticklabels(list(self.labels.values()))

        if tick_params:
            cb.ax.tick_params(**tick_params)

        return cb

    # .........................................................................

    def _parse_boundaries(self, bins):
        """Parses the boundaries for the
        :py:class:`~matplotlib.colors.BoundaryNorm`.

        Args:
            bins (Sequence): Either monotonically increasing sequence of bin
                centers or sequence of connected intervals (2-tuples).

        Returns:
            (list): Monotonically increasing boundaries.

        Raises:
            ValueError: On disconnected intervals or decreasing boundaries.
        """

        def from_intervals(intervals):
            """Extracts bin edges from sequence of connected intervals"""
            b = [intervals[0][0]]

            for low, up in intervals:
                if up < low:
                    raise ValueError(
                        "Received decreasing boundaries: " f"{up} < {low}."
                    )

                elif b[-1] != low:
                    raise ValueError(
                        "Received disconnected intervals: Upper "
                        f"bound {b[-1]} and lower bound {low} of "
                        "the proximate interval do not match."
                    )

                b.append(up)

            return b

        def from_centers(centers):
            """Calculates the bin edges as the halfway points between adjacent
            bin centers."""
            centers = np.array(centers)

            if len(centers) < 2:
                raise ValueError(
                    "At least 2 bin centers must be given to "
                    f"create a BoundaryNorm. Got: {centers}"
                )

            halves = 0.5 * np.diff(centers)

            b = (
                [centers[0] - halves[0]]
                + [c + h for c, h in zip(centers, halves)]
                + [centers[-1] + halves[-1]]
            )

            return b

        if isinstance(bins[0], tuple):
            boundaries = from_intervals(bins)
        else:
            boundaries = from_centers(bins)

        return boundaries

    def _create_cmap(
        self,
        name: str = None,
        *,
        bad: Union[str, dict] = None,
        under: Union[str, dict] = None,
        over: Union[str, dict] = None,
        **cmap_kwargs,
    ) -> mpl.colors.Colormap:
        """Creates a colormap.

        Args:
            name (str, optional): The colormap name. Can either be the name of
                a registered colormap or ``ListedColormap``. ``None`` means
                that the ``image.cmap`` value from the RC parameters is used.
            bad (Union[str, dict], optional): Set color to be used for masked
                values.
            under (Union[str, dict], optional): Set the color for low
                out-of-range values when ``norm.clip = False``.
            over (Union[str, dict], optional): Set the color for high
                out-of-range values when ``norm.clip = False``.
            **cmap_kwargs: If ``name = ListedColormap``, passed on to the
                constructor of the colormap, else passed to
                matplotlib.cm.get_cmap.

        Returns:
            matplotlib.colors.Colormap: The created colormap.

        Raises:
            ValueError: On invalid colormap name.
        """
        if name == "ListedColormap":
            cmap = mpl.colors.ListedColormap(**cmap_kwargs)
        elif name == "LinearSegmentedColormap":
            cmap = mpl.colors.LinearSegmentedColormap(name=name, **cmap_kwargs)
        else:
            try:
                cmap = copy.copy(mpl.cm.get_cmap(name, **cmap_kwargs))

            except ValueError as err:
                raise ValueError(
                    f"Received invalid colormap name: '{name}'. "
                    "Must name a registered colormap."
                ) from err

        if isinstance(bad, str):
            bad = dict(color=bad)

        if isinstance(under, str):
            under = dict(color=under)

        if isinstance(over, str):
            over = dict(color=over)

        if bad is not None:
            cmap.set_bad(**bad)

        if under is not None:
            cmap.set_under(**under)

        if over is not None:
            cmap.set_over(**over)

        return cmap

    def _create_norm(
        self, name: str = None, **norm_kwargs
    ) -> "matplotlib.colors.Normalize":
        r"""Creates a norm.

        Args:
            name (str, optional): The norm name. Must name a
                :py:class:`matplotlib.colors.Normalize` instance (see
                `matplotlib.colors <https://matplotlib.org/api/colors_api.html>`_).
                ``None`` means that the base class, *Normalize*, will be used.
            \**norm_kwargs: Passed on to the constructor of the norm.

        Returns:
            matplotlib.colors.Normalize: The created norm.

        Raises:
            ValueError: On invalid norm specification.
        """
        if name is None:
            name = "Normalize"

        if name not in NORMS:
            available_norms = ", ".join(NORMS.keys())
            raise ValueError(
                f"Received invalid norm specifier: '{name}'. "
                f"Must be one of: {available_norms}"
            )

        return NORMS[name](**norm_kwargs)


# -- Supporting functions -----------------------------------------------------


def parse_cmap_and_norm_kwargs(
    *, _key_map: dict = None, use_color_manager: bool = True, **kws
) -> dict:
    """A function that parses colormap-related keyword arguments and passes
    them through the :py:class:`.ColorManager`, making its functionality
    available in places that would otherwise not be able to use the expanded
    syntax of the color manager.

    .. note::

        The resulting dict will only have the ``cmap`` and ``cbar`` kwargs
        (or their mapped equivalents) set from the color manager, all other
        arguments are simply passed through.

        In particular, this means that the ``labels`` feature of the color
        manager is *not* supported, because this function has no ability to
        set the colorbar.

    Args:
        _key_map (dict, optional): If custom keyword argument keys are
            expected as output, e.g. ``hue_cmap`` instead of ``cmap``, set the
            values to these custom names: ``{"cmap": "hue_cmap"}``.
            Expected keys are ``cmap``, ``norm``, ``vmin``, ``vmax``. If not
            set or partially not set, will use defaults.
        use_color_manager (bool, optional): If false, will simply pass through
        **kws: Keyword arguments to parse

    No Longer Returned:
        dict: The updated keyword arguments with ``cmap`` and ``norm`` (or
            equivalent keys according to ``_key_map``).
    """
    if not use_color_manager:
        return kws

    _key = dict(cmap="cmap", norm="norm", vmin="vmin", vmax="vmax")
    if _key_map is not None:
        _key.update(_key_map)

    if _key["cmap"] not in kws and _key["norm"] not in kws:
        return kws

    # otherwise: Create a ColorManager
    cm = ColorManager(
        cmap=kws.get(_key["cmap"]),
        norm=kws.get(_key["norm"]),
        vmin=kws.get(_key["vmin"]),
        vmax=kws.get(_key["vmax"]),
    )

    # Evaluate it, only setting keys if they were there before
    if kws.get(_key["cmap"]) is not None:
        kws[_key["cmap"]] = cm.cmap
    if kws.get(_key["norm"]) is not None:
        kws[_key["norm"]] = cm.norm

    return kws
