__all__ = ["LOSS_LABELS", "plot_loss_curves"]

from collections.abc import Sequence

import matplotlib.figure

import matplotlib.pyplot as plt

import numpy as np

from move.visualization.style import (

    DEFAULT_PLOT_STYLE,

    DEFAULT_QUALITATIVE_PALETTE,

    color_cycle,

    style_settings,

)

LOSS_LABELS = ("Loss", "Cross-Entropy", "Sum of Squared Errors", "KLD")

def plot_loss_curves(

    losses: Sequence[list[float]],

    labels: Sequence[str] = LOSS_LABELS,

    style: str = DEFAULT_PLOT_STYLE,

    colormap: str = DEFAULT_QUALITATIVE_PALETTE,

) -> matplotlib.figure.Figure:

    """Plot one or more loss curves.

    Args:

        losses: List containing lists of loss values

        labels: List containing names of each loss line

        style: Name of style to apply to the plot

        colormap: Name of colormap to use for the curves

    Returns:

        Figure

    """

    num_epochs = len(losses[0])

    epochs = np.arange(num_epochs)

    with style_settings(style), color_cycle(colormap):

        fig, ax = plt.subplots()

        for loss, label in zip(losses, labels):

            ax.plot(epochs, loss, label=label, linestyle="-")

        ax.legend()

        ax.set(xlabel="Epochs", ylabel="Loss")

    return fig