"""Plotting utilities for Composer callbacks."""

from collections import defaultdict
from dataclasses import dataclass
from io import BytesIO

import numpy as np
import seaborn as sns
import torch
from matplotlib import pyplot as plt
from matplotlib.figure import Figure
from PIL import Image

import wandb


def fig_to_wandb_image(fig: Figure) -> wandb.Image:
    """Convert a matplotlib Figure to a wandb.Image.

    Parameters
    ----------
    fig: Figure
        The matplotlib Figure to convert.

    Returns
    -------
    wandb.Image
        A wandb.Image object containing the figure.

    """
    with BytesIO() as buf:
        fig.savefig(buf, format="png", dpi=150, bbox_inches="tight")
        buf.seek(0)

        pil_img = Image.open(buf)
        img = wandb.Image(pil_img)
        plt.close(fig)
        return img


@dataclass(slots=True)
class HeatmapOptions:
    """Options for configuring generic heatmap plots.

    Attributes
    ----------
    title : str | None
        Custom title for the plot. If ``None``, a default title is used.
    xlabels : list[str] | None
        Explicit x-axis tick labels. If ``None``, derived from expert_ids or
        indices.
    ylabels : list[str] | None
        Explicit y-axis tick labels. If ``None``, derived from expert_ids or
        indices.
    cmap : str
        Matplotlib/Seaborn colormap. Default is ``"viridis"``.
    annot : bool
        Whether to annotate the heatmap cells. Default is ``True``.
    fmt : str
        String formatting code for annotations. Default is ``".2f"``.
    cbar : bool
        Whether to draw a colorbar. Default is ``True``.
    square : bool
        If ``True``, set each cell to be square-shaped. Default is ``True``.
    vmin : float | None
        Minimum value for colormap normalization. Default is ``None``.
    vmax : float | None
        Maximum value for colormap normalization. Default is ``None``.
    figsize : tuple[float, float] | None
        Figure size in inches as ``(width, height)``. If ``None``, inferred from
        matrix shape.

    """

    title: str | None = None
    xlabels: list[str] | None = None
    ylabels: list[str] | None = None
    cmap: str = "viridis"
    annot: bool = True
    fmt: str = ".2f"
    cbar: bool = True
    square: bool = True
    vmin: float | None = None
    vmax: float | None = None
    figsize: tuple[float, float] | None = None


def create_heatmap(
    similarity_matrix: torch.Tensor,
    expert_ids: list[int],
    step: int,
    opts: HeatmapOptions | None = None,
) -> wandb.Image:
    """Create a generic heatmap figure from a similarity (or any) matrix.

    Parameters
    ----------
    similarity_matrix: torch.Tensor
        A square tensor representing the cosine similarity between routers.
    expert_ids: list[int]
        A list of expert IDs to use for the x-axis and y-axis labels.
    step: int
        The current step at which the heatmap is plotted.
    opts : HeatmapOptions | None, optional
        Options controlling labels, title, and styling. If ``None``, defaults are
        used.

    Returns
    -------
    wandb.Image
        A wandb.Image object containing the heatmap.

    Notes
    -----
    This function remains backward-compatible with prior usages that passed only
    ``(similarity_matrix, expert_ids, step)``. The ``opts`` parameter allows
    customizing labels, title, and styling so the function can be reused beyond
    router-weight plots (e.g., expert activation similarity heatmaps).

    """
    # Normalize options
    opts = opts or HeatmapOptions()

    # Convert to CPU numpy for plotting
    mat = similarity_matrix.detach().cpu().numpy()
    h_shape, v_shape = mat.shape

    # Determine default labels if not provided
    xlabels = opts.xlabels
    ylabels = opts.ylabels
    if xlabels is None:
        if expert_ids:
            xlabels = [f"Exp. {expert_ids[i]}" for i in range(v_shape)]
        else:
            xlabels = [str(i) for i in range(v_shape)]
    if ylabels is None:
        if expert_ids:
            ylabels = [f"Exp. {expert_ids[i]}" for i in range(h_shape)]
        else:
            ylabels = [str(i) for i in range(h_shape)]

    # Default title
    title = opts.title
    if title is None:
        base = "Similarity Heatmap"
        title = f"Step {step} - {base}"

    # Default figure size if not provided
    figsize = opts.figsize
    if figsize is None:
        figsize = (max(4, v_shape), max(3, h_shape))

    fig = plt.figure(figsize=figsize)
    sns.heatmap(
        mat,
        cmap=opts.cmap,
        annot=opts.annot,
        fmt=opts.fmt,
        cbar=opts.cbar,
        square=opts.square,
        xticklabels=xlabels,
        yticklabels=ylabels,
        vmin=opts.vmin,
        vmax=opts.vmax,
    )
    plt.title(title)
    plt.xlabel("")
    plt.ylabel("")

    return fig_to_wandb_image(fig)


def create_expert_heatmap(
    block_histograms: dict[int, torch.Tensor],
    expert_ids: list[int],
    step: int,
) -> wandb.Image:
    """Create a heatmap figure for expert selections.

    Parameters
    ----------
    block_histograms: dict[int, torch.Tensor]
        A dict of tensors containing expert selections for possible value of k in the
        topk selection.
    expert_ids: list[int]
        A list of expert IDs to use for the x-axis labels.
    step: int
        The step size for the heatmap.

    Returns
    -------
    wandb.Image
        A wandb.Image object containing the heatmap.

    """
    # Since each element in block_histograms contains the binned counts of expert
    # selections, we would like to get the normalized counts so as to have a
    # transferable colormap.
    for k, v in block_histograms.items():
        block_histograms[k] = v / v.sum()
    # Sort the block histograms by their keys (values of k for the top-k selection)
    block_histograms = dict(sorted(block_histograms.items()))
    # Stack layer's block histograms into a 2D tensor
    heatmap_data = torch.stack(list(block_histograms.values()))
    n_k, n_experts = heatmap_data.shape

    # Convert tensors to numpy arrays
    heatmap_data = heatmap_data.numpy()

    fig = plt.figure(figsize=(n_experts, n_k))
    # Main heatmap (per layer)
    sns.heatmap(
        heatmap_data,
        cmap="viridis",
        annot=True,
        cbar=True,
        square=True,
        xticklabels=[
            f"Exp. {expert_ids[i] if expert_ids else i}"
            for i in range(heatmap_data.shape[1])
        ],
        yticklabels=[f"Top-{i + 1}" for i in range(heatmap_data.shape[0])],
    )
    plt.title(f"Step {step} - Expert Selection Heatmap")
    plt.xlabel("")
    plt.ylabel("")

    return fig_to_wandb_image(fig)


def create_thresholded_expert_scores_histogram(
    block_token_wise_threshold_histogram: torch.Tensor,
    step: int,
) -> wandb.Image:
    """Create a histogram figure for expert thresholded scores.

    Parameters
    ----------
    block_token_wise_threshold_histogram: torch.Tensor
        A tensor containing histogram counts with the token-wise thresholded scores for
        each expert.
    step: int
        The current training/evaluation step number for labeling the plot.

    Returns
    -------
    wandb.Image
        A wandb.Image object containing the histogram.

    """
    # Create the figure and axis
    fig, ax = plt.subplots(figsize=(6, 6))
    # Get the number of bins from the tensor
    n_bins = block_token_wise_threshold_histogram.shape[0]
    # Create a simple histogram from the binned block_token_wise_threshold_histogram
    # whose shape is (n_bins,) containing the counts
    ax.bar(
        x=range(n_bins),
        height=block_token_wise_threshold_histogram.numpy(),
        tick_label=[f"{i}" for i in range(n_bins)],
    )
    ax.set_title(f"Step {step} - Expert Thresholded Scores Histogram")
    ax.set_xlabel("# Selected Experts w/ threshold")
    ax.set_ylabel("Token count")
    # Return the image
    return fig_to_wandb_image(fig)


def create_expert_scores_histogram(
    block_scores_histograms: dict[int, tuple[torch.Tensor, torch.Tensor]],
    expert_ids: list[int],
    step: int,
) -> tuple[wandb.Image, dict[int, dict[str, float]]]:
    """Create a histogram figure for expert selection scores.

    Parameters
    ----------
    block_scores_histograms: dict[int, tuple[torch.Tensor, torch.Tensor]]
        A dict of tuples containing histograms of scores and the corresponding bin edges
        for each expert.
    expert_ids: list[int]
        A list of expert IDs to use for the x-axis labels.
    step: int
        The step size for the histogram.

    Returns
    -------
    tuple[wandb.Image, dict[int, dict[str, float]]]
        A tuple containing a wandb.Image object with the histogram and a
        dictionary with statistics (mean, median, std, lower quartile, upper quartile)
        for each expert. The dictionary keys are the expert IDs or indices if
        `expert_ids` is empty.

    """
    # Initialize the dictionary with the statistics to log
    ret_statistics: dict[int, dict[str, float]] = defaultdict(dict)
    # Get the number of experts from the histograms
    n_experts = len(block_scores_histograms)
    fig, ax = plt.subplots(figsize=(6, 6))

    # Create a color palette for the experts
    colors = plt.get_cmap("tab10")(range(n_experts))

    # Loop over the histograms and plot them as line plots with vertical stacking
    for i, (bin_counts, bin_edges) in block_scores_histograms.items():
        # Calculate the cumulative
        cumulative_bin_counts = bin_counts.cumsum(0).numpy()
        # Normalize the bin counts to ensure they sum to 1
        cumulative_bin_counts /= cumulative_bin_counts.max()
        # Find the median from the cumulative
        median_index = np.searchsorted(cumulative_bin_counts, 0.5)
        # Find the 25th and 75th quantiles from the cumulative
        lower_quartile_index = np.searchsorted(cumulative_bin_counts, 0.25)
        upper_quartile_index = np.searchsorted(cumulative_bin_counts, 0.75)
        # Calculate bin centers for line plot
        bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
        # Calculate the mean of the histogram
        bin_centers_mean = (
            bin_centers.numpy() * bin_counts.numpy()
        ).sum() / bin_counts.numpy().sum()
        # Calculate the standard deviation of the histogram
        bin_centers_std = (
            (bin_counts.numpy() * (bin_centers.numpy() - bin_centers_mean) ** 2).sum()
            / bin_counts.numpy().sum()
        ) ** 0.5
        ret_statistics[expert_ids[i] if expert_ids else i].update(
            {
                "mean": bin_centers_mean,
                "median": bin_centers.numpy()[median_index],
                "std": bin_centers_std,
                "lower_quartile": bin_centers.numpy()[lower_quartile_index],
                "upper_quartile": bin_centers.numpy()[upper_quartile_index],
            },
        )
        # Plot the histogram as a line plot
        ax.plot(
            bin_centers.numpy(),
            bin_counts.numpy() / bin_counts.numpy().sum(),
            label=f"Expert-{expert_ids[i] if expert_ids else i}",
            linewidth=2,
            color=colors[i % len(colors)],
        )
        # Add mean, median and std to the plot
        ax.axvline(
            x=bin_centers_mean,
            color=colors[i % len(colors)],
            linestyle="--",
        )
        ax.axvline(
            x=bin_centers.numpy()[median_index],
            color=colors[i % len(colors)],
            linestyle=":",
        )
        ax.axvline(
            x=bin_centers_mean + bin_centers_std,
            color=colors[i % len(colors)],
            linestyle="-.",
        )
        ax.axvline(
            x=bin_centers_mean - bin_centers_std,
            color=colors[i % len(colors)],
            linestyle="-.",
        )

    # TODO(anon): Fit the histograms with a beta distribution. We believe the target
    # distribution these routers should point to for balanced training is a beta
    # distribution with alpha=beta < 1.0.

    # Set the x-lim to match the score range
    ax.set_xlim((0, 1.0))
    ax.set_title(f"Step {step} - Expert Selection Scores Distribution")
    ax.set_xlabel("Scores (sigmoided)")
    ax.set_ylabel("Tokens")
    ax.grid(visible=True, alpha=0.3)

    # Add legend entries for statistics line styles (in gray)
    ax.axvline(x=-1, color="gray", linestyle="--", label="Mean")
    ax.axvline(x=-1, color="gray", linestyle=":", label="Median")
    ax.axvline(x=-1, color="gray", linestyle="-.", label="±1 Std")

    # Add legend for mean, median and std lines
    ax.legend()
    return fig_to_wandb_image(fig), ret_statistics


def create_expert_tokens_scores_histogram(
    block_full_selection_history: torch.Tensor,
    expert_ids: list[int],
    step: int,
) -> wandb.Image:
    """Create a vertically stacked histogram of expert token scores across time.

    This function creates a visualization showing how expert selection scores
    vary across tokens, with each expert's scores displayed as a separate
    distribution stacked vertically. All experts share the same x-axis (token
    positions) but are offset vertically with visual separators between them.

    Parameters
    ----------
    block_full_selection_history: torch.Tensor
        A tensor containing the full selection history for experts, with shape
        (n_tokens, n_experts). Each element represents the selection score
        for a specific expert at a specific token position.
    expert_ids: list[int]
        A list of expert IDs to use for labeling. If empty, expert indices
        will be used directly.
    step: int
        The step size for the histogram.

    Returns
    -------
    wandb.Image
        A wandb.Image object containing the vertically stacked plot
        showing expert token scores distribution over time, with each expert
        occupying its own vertical band.

    """
    # Get the number of experts, sequence length, and batch size
    batch_size = block_full_selection_history.shape[0]
    sequence_length = block_full_selection_history.shape[1]
    n_experts = block_full_selection_history.shape[2]

    fig, ax = plt.subplots(figsize=(2 * batch_size, 2 * n_experts))
    # Create a color palette for the experts
    colors = plt.get_cmap("tab10")(range(n_experts))

    # Create vertically stacked distributions
    base_y_offset = 1.1
    for k in range(n_experts):
        # Create a new x-axis (horizontal) for each expert to create a vertical stacking
        # effect across experts
        ax_k = ax.twiny() if k > 0 else ax
        ax_k.spines["top"].set_visible(False)
        ax_k.spines["right"].set_visible(False)
        ax_k.xaxis.set_ticks_position("bottom")
        ax_k.yaxis.set_ticks_position("right")
        # Shift the y=0 horizontal line value up by k units to create vertical stacking
        y_offset = k * base_y_offset
        ax_k.spines["bottom"].set_position(("data", y_offset))

        # Loop through each sample in the batch plotting the expert scores
        for s in range(batch_size):
            # Create a horizontal offset for each sample in the batch
            x_offset = s * sequence_length
            sample_expert_scores = block_full_selection_history[s, :, k].numpy()

            # Use the same x-axis (token positions) for all experts
            x_axis = torch.arange(sequence_length) + x_offset

            # Apply vertical offset to stack distributions vertically
            y_values = sample_expert_scores + y_offset
            ax_k.plot(
                x_axis,
                y_values,
                linewidth=1.0,
                color=colors[k % len(colors)],
                alpha=0.8,
            )
            ax_k.boxplot(
                y_values,
                notch=True,
                sym="",
                orientation="vertical",
                positions=[x_offset + sequence_length / 2],
                widths=sequence_length / 2,
                tick_labels=[""],
                showmeans=True,
                meanline=True,
                showcaps=True,
                showbox=True,
            )

        # Remove the horizontal labels when k > 0
        if k > 0:
            ax_k.xaxis.set_ticklabels([])

        # Change the x-ticks to be multiples of the sequence length, and set the x-lim
        # to match the last sequence
        ax_k.set_xticks(
            ticks=[(s + 0.5) * sequence_length for s in range(batch_size)],
            labels=[f"Seq. {s}" for s in range(1, batch_size + 1)],
        )
        ax_k.set_xlim((0, batch_size * sequence_length))

    # Adding a vertical line at the last x position of every sample
    for s in range(batch_size + 1):
        x_offset = s * sequence_length
        ax.axvline(
            x=x_offset,
            color="black",
            linestyle="--",
            alpha=0.5,
            linewidth=0.5,
        )

    # Set the vertical ticks to correspond to the 0 and 1 of each expert
    y_ticks = (
        [0 + k * base_y_offset for k in range(n_experts)]
        + [1 + k * base_y_offset for k in range(n_experts)]
        + [0.5 + k * base_y_offset for k in range(n_experts)]
    )
    y_ticks.sort()
    ax.set_yticks(ticks=y_ticks, labels=["0", "0.5", "1"] * n_experts)

    # Customize the plot
    ax.set_title(
        f"Step {step} - Expert Token Scores Distribution",
        fontsize=14,
        fontweight="bold",
    )
    ax.set_xlabel("Token Position", fontsize=12)
    ax.set_ylabel("Expert Selection Scores", fontsize=12)
    ax.grid(visible=True, alpha=0.3)

    # Set y-tick labels to show expert boundaries

    # Add secondary y-axis labels for expert identification
    ax2 = ax.twinx()
    ax2.set_ylim(ax.get_ylim())
    ax2.set_yticks([0.5 + k * base_y_offset for k in range(n_experts)])
    ax2.set_yticklabels(
        [f"Expert-{expert_ids[k] if expert_ids else k}" for k in range(n_experts)],
        fontsize=10,
    )
    ax2.set_ylabel("Expert ID", fontsize=12)

    # Adjust layout to prevent legend cutoff
    plt.tight_layout()

    return fig_to_wandb_image(fig)


def _get_list_of_probability_distribution_references(
    experts_ids: list[int],
) -> dict[str, torch.Tensor]:
    # Initialize return dictionary with the uniform and zero distributions
    prob_dist_refs: dict[str, torch.Tensor] = {
        "uniform": torch.ones((len(experts_ids),)) / len(experts_ids),
        # Add a small value to avoid division by zero
        "zero": torch.zeros((len(experts_ids),)) + 1e-8,
    }
    # Add one peaked distribution for each expert
    for idx, expert_id in enumerate(experts_ids):
        # Create a one-hot distribution peaked at this expert
        peaked_dist = torch.zeros((len(experts_ids),)) + 1e-8
        peaked_dist[idx] = 1.0
        prob_dist_refs[f"peaked_at_{expert_id}"] = peaked_dist
    return prob_dist_refs


# NOTE: Currently unused, keeping for reference
def _plot_l1_distance_histo(  # pyright: ignore[reportUnusedFunction]
    input_tokenwise_distribution: torch.Tensor,
    target_experts_distribution: torch.Tensor,
    step: int,
    ref_name: str,
) -> wandb.Image:
    """Create a histogram of the L1 distance between token distributions.

    This function computes the L1 distance between the input token-wise distribution
    and the target expert distribution, and creates a histogram of the L1 distance
    values.

    Parameters
    ----------
    input_tokenwise_distribution: torch.Tensor
        The input token-wise distribution with shape [n_tokens, experts].
    target_experts_distribution: torch.Tensor
        The target expert distribution with shape [experts].
    step: int
        The current training step.
    ref_name: str
        The name of the reference distribution.

    Returns
    -------
    wandb.Image
        A wandb.Image object containing the histogram of L1 distances.

    """
    # Compute the L1 distance for each token
    l1_norm = (
        input_tokenwise_distribution
        - target_experts_distribution.repeat(
            input_tokenwise_distribution.shape[0],
            1,
        )
    ).abs()
    # Average the L1 distance across experts
    l1_norm_mean = l1_norm.mean(dim=-1)
    # Create a histogram of the L1 distance values
    fig, ax = plt.subplots(figsize=(10, 5))
    sns.histplot(l1_norm_mean.numpy(), kde=True, ax=ax)
    ax.set_title(f"L1 Histogram - Step {step} - {ref_name}")
    ax.set_xlabel("L1 Norm (averaged across experts)")
    ax.set_ylabel("Frequency")
    return fig_to_wandb_image(fig)


# NOTE: Currently unused, keeping for reference
def _plot_kl_distance_histo(  # pyright: ignore[reportUnusedFunction]
    input_tokenwise_distribution: torch.Tensor,
    target_experts_distribution: torch.Tensor,
    step: int,
    ref_name: str,
) -> wandb.Image:
    """Create a histogram of the KL divergence between token distributions.

    This function computes the KL divergence between the input token-wise distribution
    and the target expert distribution, and creates a histogram of the KL divergence
    values.

    Parameters
    ----------
    input_tokenwise_distribution: torch.Tensor
        The input token-wise distribution with shape [n_tokens, experts].
    target_experts_distribution: torch.Tensor
        The target expert distribution with shape [experts].
    step: int
        The current training step.
    ref_name: str
        The name of the reference distribution.

    Returns
    -------
    wandb.Image
        A wandb.Image object containing the histogram of KL divergence values.

    """
    # Compute the KL divergence for each token
    kl_div = torch.nn.functional.kl_div(
        input_tokenwise_distribution,
        target_experts_distribution.repeat(
            input_tokenwise_distribution.shape[0],
            1,
        ),
        reduction="none",
    )
    # Sum the KL divergence across experts
    kl_div_sum = kl_div.sum(dim=-1)
    # Create a histogram of the KL divergence values
    fig, ax = plt.subplots(figsize=(10, 5))
    sns.histplot(kl_div_sum.numpy(), kde=True, ax=ax)
    ax.set_title(f"KL Divergence Histogram - Step {step} - {ref_name}")
    ax.set_xlabel("KL Divergence (summed across experts)")
    ax.set_ylabel("Frequency")
    return fig_to_wandb_image(fig)


def _plot_cos_sim_distance_histo(
    input_tokenwise_distribution: torch.Tensor,
    target_experts_distribution: torch.Tensor,
    step: int,
    ref_name: str,
) -> wandb.Image:
    """Create a histogram of the cosine similarity between token distributions.

    This function computes the cosine similarity between the input token-wise
    distribution and the target expert distribution, and creates a histogram of the
    cosine similarity values.

    Parameters
    ----------
    input_tokenwise_distribution: torch.Tensor
        The input token-wise distribution with shape [n_tokens, experts].
    target_experts_distribution: torch.Tensor
        The target expert distribution with shape [experts].
    step: int
        The current training step.
    ref_name: str
        The name of the reference distribution.

    Returns
    -------
    wandb.Image
        A wandb.Image object containing the histogram of cosine similarity values.

    """
    # Compute the cosine similarity for each token
    cos_sim = torch.nn.functional.cosine_similarity(
        input_tokenwise_distribution,
        target_experts_distribution.repeat(
            input_tokenwise_distribution.shape[0],
            1,
        ),
        dim=-1,
    )
    # Create a histogram of the cosine similarity values
    fig, ax = plt.subplots(figsize=(10, 5))
    sns.histplot(cos_sim.numpy(), kde=True, ax=ax)
    ax.set_title(f"Cosine Similarity Histogram - Step {step} - {ref_name}")
    ax.set_xlabel("Cosine Similarity (averaged across experts)")
    ax.set_ylabel("Frequency")
    return fig_to_wandb_image(fig)


def create_tokens_kl_histograms(
    block_full_selection_history: torch.Tensor,
    expert_ids: list[int],
    step: int,
    max_entries: int = 10_000,
) -> dict[str, wandb.Image]:
    """Create L1 distance histograms comparing token probabilities to references.

    Parameters
    ----------
    block_full_selection_history: torch.Tensor
        A tensor containing the full selection history for experts, with shape
        (batch_size, sequence_length, n_experts). Each element represents the
        selection score for a specific expert at a specific token position.
    expert_ids: list[int]
        A list of expert IDs to use for labeling. If empty, expert indices
        will be used directly.
    step: int
        The step at which the histogram is plotted.
    max_entries: int
        The maximum number of entries to include in the histograms. Defaults to 10,000.

    Returns
    -------
    dict[str, wandb.Image]
        A dictionary mapping distribution names to wandb.Image objects
        containing histograms of L1 distances between per-token probabilities
        and reference distributions.

    """
    # Initialize the return dictionary
    figures_dict: dict[str, wandb.Image] = {}
    # NOTE: The block_full_selection_history is a tensor of shape
    # [batch_size, sequence_length, n_experts] containing the sigmoided
    # scores from the router for each token seen during the current batch
    # Get the per-token probability scores
    per_token_scores = block_full_selection_history.view(
        -1,
        block_full_selection_history.shape[-1],
    )
    # Randomly slice `max_entries` samples from the `per_token_scores`
    indices = torch.randperm(
        per_token_scores.shape[0],
    )[:max_entries]
    per_token_scores = per_token_scores[indices]
    per_token_probabilities = (
        per_token_scores / per_token_scores.sum(dim=-1, keepdim=True)
    ) + 1e-8  # Add a small value to avoid division by zero

    # If expert_ids is empty, create a fallback list using indices
    if not expert_ids:
        n_experts = per_token_probabilities.shape[-1]
        expert_ids = list(range(n_experts))

    # Get the reference distributions
    prob_dist_refs = _get_list_of_probability_distribution_references(
        experts_ids=expert_ids,
    )
    # Loop over the reference distributions and compute the CosSim distance
    for ref_name, ref_distribution in prob_dist_refs.items():
        # Only computing CosSim distances histograms
        figures_dict[f"cos_sim_histogram_{ref_name}"] = _plot_cos_sim_distance_histo(
            input_tokenwise_distribution=per_token_probabilities,
            target_experts_distribution=ref_distribution,
            step=step,
            ref_name=ref_name,
        )
    return figures_dict
