"""Provides grid_archive_heatmap."""

from __future__ import annotations

from collections.abc import Sequence
from typing import Literal

import matplotlib.colors
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from matplotlib.typing import ColorType
from pandas import DataFrame

from ribs.archives import ArchiveDataFrame, GridArchive
from ribs.visualize._utils import (
    archive_heatmap_1d,
    retrieve_cmap,
    set_cbar,
    validate_df,
    validate_heatmap_visual_args,
)


def grid_archive_heatmap(
    archive: GridArchive,
    ax: Axes | None = None,
    *,
    df: DataFrame | ArchiveDataFrame | None = None,
    transpose_measures: bool = False,
    cmap: str | Sequence[ColorType] | matplotlib.colors.Colormap = "magma",
    aspect: Literal["auto", "equal"] | float | None = None,
    vmin: float | None = None,
    vmax: float | None = None,
    cbar: Literal["auto"] | None | Axes = "auto",
    cbar_kwargs: dict | None = None,
    rasterized: bool = False,
    pcm_kwargs: dict | None = None,
) -> None:
    r"""Plots heatmap of a :class:`~ribs.archives.GridArchive` with 1D or 2D measure space.

    This function creates a grid of cells and shades each cell with a color
    corresponding to the objective value of that cell's elite. This function uses
    :func:`~matplotlib.pyplot.pcolormesh` to generate the grid. For further
    customization, pass extra kwargs to :func:`~matplotlib.pyplot.pcolormesh` through
    the ``pcm_kwargs`` parameter. For instance, to create black boundaries of width 0.1,
    pass in ``pcm_kwargs={"edgecolor": "black", "linewidth": 0.1}``.

    Examples:
        .. plot::
            :context: close-figs

            Heatmap of a 2D GridArchive

            >>> import numpy as np
            >>> import matplotlib.pyplot as plt
            >>> from ribs.archives import GridArchive
            >>> from ribs.visualize import grid_archive_heatmap
            >>> # Populate the archive with the negative sphere function.
            >>> archive = GridArchive(solution_dim=2,
            ...                       dims=[20, 20],
            ...                       ranges=[(-1, 1), (-1, 1)])
            >>> x = np.random.uniform(-1, 1, 10000)
            >>> y = np.random.uniform(-1, 1, 10000)
            >>> archive.add(solution=np.stack((x, y), axis=1),
            ...             objective=-(x**2 + y**2),
            ...             measures=np.stack((x, y), axis=1))
            >>> # Plot a heatmap of the archive.
            >>> plt.figure(figsize=(8, 6))
            >>> grid_archive_heatmap(archive)
            >>> plt.title("Negative sphere function")
            >>> plt.xlabel("x coords")
            >>> plt.ylabel("y coords")
            >>> plt.show()

        .. plot::
            :context: close-figs

            Heatmap of a 1D GridArchive

            >>> import numpy as np
            >>> import matplotlib.pyplot as plt
            >>> from ribs.archives import GridArchive
            >>> from ribs.visualize import grid_archive_heatmap
            >>> # Populate the archive with the negative sphere function.
            >>> archive = GridArchive(solution_dim=2,
            ...                       dims=[20], ranges=[(-1, 1)])
            >>> x = np.random.uniform(-1, 1, 1000)
            >>> archive.add(solution=np.stack((x, x), axis=1),
            ...             objective=-x**2,
            ...             measures=x[:, None])
            >>> # Plot a heatmap of the archive.
            >>> plt.figure(figsize=(8, 6))
            >>> grid_archive_heatmap(archive)
            >>> plt.title("Negative sphere function with 1D measures")
            >>> plt.xlabel("x coords")
            >>> plt.show()

    Args:
        archive: A 1D or 2D :class:`~ribs.archives.GridArchive`.
        ax: Axes on which to plot the heatmap.  If ``None``, the current axis will be
            used.
        df: If provided, we will plot data from this argument instead of the data
            currently in the archive. This data can be obtained by, for instance,
            calling :meth:`ribs.archives.ArchiveBase.data` with ``return_type="pandas"``
            and modifying the resulting :class:`~ribs.archives.ArchiveDataFrame`. Note
            that, at a minimum, the data must contain columns for index, objective, and
            measures. To display a custom metric, replace the "objective" column.
        transpose_measures: By default, the first measure in the archive will appear
            along the x-axis, and the second will be along the y-axis. To switch this
            behavior (i.e. to transpose the axes), set this to ``True``. Does not apply
            for 1D archives.
        cmap: The colormap to use when plotting intensity. Either the name of a
            :class:`~matplotlib.colors.Colormap`, a list of Matplotlib color
            specifications (e.g., an :math:`N \times 3` or :math:`N \times 4` array --
            see :class:`~matplotlib.colors.ListedColormap`), or a
            :class:`~matplotlib.colors.Colormap` object.
        aspect: The aspect ratio of the heatmap (i.e. height/width). Defaults to
            ``'auto'`` for 2D and ``0.5`` for 1D. ``'equal'`` is the same as
            ``aspect=1``. See :meth:`matplotlib.axes.Axes.set_aspect` for more info.
        vmin: Minimum objective value to use in the plot. If ``None``, the minimum
            objective value in the archive is used.
        vmax: Maximum objective value to use in the plot. If ``None``, the maximum
            objective value in the archive is used.
        cbar: By default, this is set to ``'auto'`` which displays the colorbar on the
            archive's current :class:`~matplotlib.axes.Axes`. If ``None``, then colorbar
            is not displayed. If this is an :class:`~matplotlib.axes.Axes`, displays the
            colorbar on the specified Axes.
        cbar_kwargs: Additional kwargs to pass to :func:`~matplotlib.pyplot.colorbar`.
        rasterized: Whether to rasterize the heatmap. This can be useful for saving to a
            vector format like PDF. Essentially, only the heatmap will be converted to a
            raster graphic so that the archive cells will not have to be individually
            rendered. Meanwhile, the surrounding axes, particularly text labels, will
            remain in vector format.
        pcm_kwargs: Additional kwargs to pass to :func:`~matplotlib.pyplot.pcolormesh`.

    Raises:
        ValueError: The archive's measure dimension must be 1D or 2D.
    """
    validate_heatmap_visual_args(
        aspect,
        cbar,
        archive.measure_dim,
        [1, 2],
        "Heatmap can only be plotted for a 1D or 2D GridArchive",
    )

    if aspect is None:
        # Handles default aspects for different dims.
        if archive.measure_dim == 1:
            aspect = 0.5
        else:
            aspect = "auto"

    # Try getting the colormap early in case it fails.
    cmap = retrieve_cmap(cmap)

    # Retrieve archive data.
    if df is None:
        index_batch = archive.data("index")
        objective_batch = archive.data("objective")
    else:
        df = validate_df(df)
        index_batch = df["index"]
        objective_batch = df["objective"]

    if archive.measure_dim == 1:
        cell_objectives = np.full(archive.cells, np.nan)
        cell_idx = archive.int_to_grid_index(index_batch).squeeze()
        cell_objectives[cell_idx] = objective_batch

        archive_heatmap_1d(
            archive,
            cell_boundaries=archive.boundaries[0],
            cell_objectives=cell_objectives,
            ax=ax,
            cmap=cmap,
            aspect=aspect,
            vmin=vmin,
            vmax=vmax,
            cbar=cbar,
            cbar_kwargs=cbar_kwargs,
            rasterized=rasterized,
            pcm_kwargs=pcm_kwargs,
        )

    elif archive.measure_dim == 2:
        # Retrieve data from archive.
        lower_bounds = archive.lower_bounds
        upper_bounds = archive.upper_bounds
        x_dim, y_dim = archive.dims
        x_bounds = archive.boundaries[0]
        y_bounds = archive.boundaries[1]

        # Color for each cell in the heatmap.
        colors = np.full((y_dim, x_dim), np.nan)
        grid_index_batch = archive.int_to_grid_index(index_batch)
        colors[grid_index_batch[:, 1], grid_index_batch[:, 0]] = objective_batch

        if transpose_measures:
            # Since the archive is 2D, transpose by swapping the x and y boundaries and
            # by flipping the bounds (the bounds are arrays of length 2).
            x_bounds, y_bounds = y_bounds, x_bounds
            lower_bounds = np.flip(lower_bounds)
            upper_bounds = np.flip(upper_bounds)
            colors = colors.T

        # Initialize the axis.
        ax = plt.gca() if ax is None else ax
        ax.set_xlim(lower_bounds[0], upper_bounds[0])  # ty: ignore[invalid-argument-type]
        ax.set_ylim(lower_bounds[1], upper_bounds[1])  # ty: ignore[invalid-argument-type]

        ax.set_aspect(aspect)

        # Create the plot.
        pcm_kwargs = {} if pcm_kwargs is None else pcm_kwargs
        vmin = (
            np.min(objective_batch)
            if vmin is None and len(objective_batch) > 0
            else vmin
        )
        vmax = (
            np.max(objective_batch)
            if vmax is None and len(objective_batch) > 0
            else vmax
        )
        t = ax.pcolormesh(
            x_bounds,
            y_bounds,
            colors,
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
            rasterized=rasterized,
            **pcm_kwargs,
        )

        # Create color bar.
        set_cbar(t, ax, cbar, cbar_kwargs)
