"""Function to make a categorical scatter plot."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import matplotlib.pyplot as plt


def categorical_scatter_2d(
    X2D,
    class_idxs,
    ms=3,
    ax=None,
    alpha=0.1,
    legend=True,
    figsize=None,
    show=False,
    savename=None,
):
    """Plot a 2D scatter plot with markers for classes."""
    ## Plot a 2D matrix with corresponding class labels: each class diff colour
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    # ax.margins(0.05) # Optional, just adds 5% padding to the autoscaling
    classes = list(np.unique(class_idxs))
    markers = "os" * len(classes)
    colors = plt.cm.rainbow(np.linspace(0, 1, len(classes)))

    for i, cls in enumerate(classes):
        mark = markers[i]
        ax.plot(
            X2D[class_idxs == cls, 0],
            X2D[class_idxs == cls, 1],
            marker=mark,
            linestyle="",
            ms=ms,
            label=str(cls),
            alpha=alpha,
            color=colors[i],
            markeredgecolor="black",
            markeredgewidth=0.4,
        )
    if legend:
        ax.legend()

    if savename is not None:
        plt.tight_layout()
        plt.savefig(savename)

    if show:
        plt.show()

    return ax
