# Licensed under a 3-clause BSD style license - see LICENSE.rst
import numpy as np
import matplotlib.pyplot as plt

# taken from the matploltlib documentation
# https://matplotlib.org/3.1.0/gallery/images_contours_and_fields/image_annotated_heatmap.html#sphx-glr-gallery-images-contours-and-fields-image-annotated-heatmap-py

__all__ = [
    "annotate_heatmap",
    "plot_heatmap",
]


def plot_heatmap(
    data, row_labels, col_labels, ax=None, cbar_kw=None, cbarlabel="", **kwargs
):
    """
    Create a heatmap from a numpy array and two lists of labels.

    Parameters
    ----------
    data : `~numpy.ndarray`
        Data array.
    row_labels : list or `~numpy.ndarray`
        List or array of labels for the rows.
    col_labels : list or `~numpy.ndarray`
        List or array of labels for the columns.
    ax : `matplotlib.axes.Axes`, optional
        Axis instance to which the heatmap is plotted. Default is None. If None, the current one is used.
    cbar_kw : dict, optional
        A dictionary with arguments to `matplotlib.Figure.colorbar`. Default is None.
    cbarlabel : str, optional
        The label for the color bar. Default is "".
    **kwargs : dict, optional
        Other keyword arguments forwarded to `matplotlib.axes.Axes.imshow`.
    """
    if ax is None:
        ax = plt.gca()

    if cbar_kw is None:
        cbar_kw = {}

    # Plot the heatmap
    im = ax.imshow(data, **kwargs)

    # Create colorbar
    cbar = ax.figure.colorbar(im, ax=ax, **cbar_kw)
    cbar.ax.set_ylabel(cbarlabel, rotation=-90, va="bottom")

    # We want to show all ticks...
    ax.set_xticks(np.arange(data.shape[1]))
    ax.set_yticks(np.arange(data.shape[0]))
    # ... and label them with the respective list entries.
    ax.set_xticklabels(col_labels)
    ax.set_yticklabels(row_labels)

    # Let the horizontal axes labeling appear on top.
    ax.tick_params(top=True, bottom=False, labeltop=True, labelbottom=False)

    # Rotate the tick labels and set their alignment.
    plt.setp(ax.get_xticklabels(), rotation=-30, ha="right", rotation_mode="anchor")

    # Turn spines off and create white grid.
    for edge, spine in ax.spines.items():
        spine.set_visible(False)

    ax.set_xticks(np.arange(data.shape[1] + 1) - 0.5, minor=True)
    ax.set_yticks(np.arange(data.shape[0] + 1) - 0.5, minor=True)
    ax.grid(which="minor", color="w", linestyle="-", linewidth=1.5)
    ax.tick_params(which="minor", bottom=False, left=False)

    return im, cbar


def annotate_heatmap(
    im,
    data=None,
    valfmt="{x:.2f}",
    textcolors=("black", "white"),
    threshold=None,
    **textkw,
):
    """
    A function to annotate a heatmap.

    Parameters
    ----------
    im
        The AxesImage to be labeled.
    data : `~numpy.ndarray`, optional
        Data used to annotate. Default is None. If None, the image's data is used.
    valfmt : str format or `matplotlib.ticker.Formatter`, optional
        The format of the annotations inside the heatmap. This should either
        use the string format method, e.g. "$ {x:.2f}"
        or be a `matplotlib.ticker.Formatter` instance. Default is "{x:.2f}".
    textcolors : list or `~numpy.ndarray`, optional
        Two color specifications.  The first is used for
        values below a threshold, the second for those above. Default is ["black", "white"].
    threshold : float, optional
        Value in data units according to which the colors from textcolors are
        applied. Default is None. If None the middle of the colormap is used as
        separation.
    **kwargs : dict, optional
        Other keyword arguments forwarded to each call to `text` used to create
        the text labels.
    """
    import matplotlib

    if not isinstance(data, (list, np.ndarray)):
        data = im.get_array()

    # Normalize the threshold to the images color range.
    if threshold is not None:
        threshold = im.norm(threshold)
    else:
        threshold = im.norm(data.max()) / 2.0

    # Set default alignment to center, but allow it to be
    # overwritten by textkw.
    kw = dict(horizontalalignment="center", verticalalignment="center")
    kw.update(textkw)

    # Get the formatter in case a string is supplied
    if isinstance(valfmt, str):
        valfmt = matplotlib.ticker.StrMethodFormatter(valfmt)

    # Loop over the data and create a `Text` for each "pixel".
    # Change the text's color depending on the data.
    texts = []
    for i in range(data.shape[0]):
        for j in range(data.shape[1]):
            kw.update(color=textcolors[int(im.norm(data[i, j]) > threshold)])
            text = im.axes.text(j, i, valfmt(data[i, j], None), **kw)
            texts.append(text)

    return texts
