""" Define color-related functions. """

import matplotlib.colors as mplc

from mpl_toolkits.axes_grid1 import make_axes_locatable


def colorbar(
    mappable,
    fraction: float = 0.05,
    pad: float = 0.05,
    location: str = "right",
    **kwargs,
):
    """ Add a colorbar to a plot so that its height matches the plot height.

    This splits off a small section of the original axes in which the mappable was drawn
    and uses it for the colorbar.

    The code here is loosely based on this Stackoverflow answer:

        https://stackoverflow.com/a/18195921

    Parameters
    ----------
    mappable
        Output from a function like `imshow`, `scatter`, or `contourf` that uses a
        color map.
    fraction
        Fraction of axes to be taken up by the colorbar.
    pad
        Amount of padding around colorbar.
    location
        Location for the colorbar ("left", "right", "top", or "bottom").
    Additional arguments are passed to `plt.Figure.colorbar`.

    Returns the colorbar generated by Matplotlib.
    """
    # find the axes
    try:
        ax = mappable.axes
    except AttributeError:
        # output of contourf doesn't have `axes` member
        ax = mappable.ax

    # find the corresponding figure
    fig = ax.figure

    # split the axes in two
    divider = make_axes_locatable(ax)
    size = f"{int(100 * fraction)}%"
    cax = divider.append_axes(location, size=size, pad=pad)

    # make and return the colorbar
    return fig.colorbar(mappable, cax=cax, **kwargs)


def gradient_cmap(name: str, c1, c2, **kwargs) -> mplc.Colormap:
    """ Make a colormap that interpolates between two given colors.

    Parameters
    ----------
    name
        Name to give the colormap.
    c1
        Starting color. This can be in any format recognized by Matplotlib.
    c2
        Final color. This can be in any format recognized by Matplotlib.

    Returns the `Colormap` object.
    """
    rgba1 = mplc.to_rgba(c1)
    rgba2 = mplc.to_rgba(c2)

    segments = {
        "red": [[0.0, 0.0, rgba1[0]], [1.0, rgba2[0], 0.0]],
        "green": [[0.0, 0.0, rgba1[1]], [1.0, rgba2[1], 0.0]],
        "blue": [[0.0, 0.0, rgba1[2]], [1.0, rgba2[2], 0.0]],
        "alpha": [[0.0, 0.0, rgba1[3]], [1.0, rgba2[3], 0.0]],
    }

    cmap = mplc.LinearSegmentedColormap(name, segmentdata=segments, **kwargs)
    return cmap
