"""
Loss functions and regularisers for SuperCluster.
"""

from __future__ import annotations

import torch
import torch.nn.functional as F


def center_diversity_loss(centers: torch.Tensor) -> torch.Tensor:
    """Prototype center diversity regulariser.

    Computes the mean pairwise cosine similarity of the normalised prototype
    centers. Minimising this loss encourages the K centers to be spread apart
    on the unit hypersphere, preventing prototype collapse without distorting
    the prediction scores.

    This loss acts only on the routing geometry (``model.centers``); it has no
    direct effect on ``model.prototype_scores``.

    Parameters
    ----------
    centers : torch.Tensor, shape [K, d]
        Raw (unnormalised) prototype center embeddings (``model.centers``).

    Returns
    -------
    torch.Tensor
        Scalar loss (mean off-diagonal cosine similarity, in [-1, 1]).

    Examples
    --------
    ::

        from supercluster import SuperCluster, center_diversity_loss
        import torch.nn as nn

        model = SuperCluster(input_dim=10, embedding_dim=64, num_clusters=8)
        logits, cw = model(x)
        loss = nn.BCEWithLogitsLoss()(logits, y)
        loss = loss + 0.1 * center_diversity_loss(model.centers)
        loss.backward()
    """
    K = centers.size(0)
    c = F.normalize(centers, dim=-1)              # [K, d]
    sim = c @ c.T                                  # [K, K]
    mask = ~torch.eye(K, dtype=torch.bool, device=centers.device)
    return sim[mask].mean()


def effective_prototype_count(
    cluster_weights: torch.Tensor, eps: float = 1e-9
) -> float:
    """Entropy-based effective prototype count K_eff.

    K_eff = exp(H(π̄)) where π̄ = mean_i π_i is the mean routing distribution.
    Equals K when routing is perfectly uniform; equals 1 when a single prototype
    captures all examples.

    Parameters
    ----------
    cluster_weights : torch.Tensor, shape [N, K]
        Soft assignment weights from a forward pass.
    eps : float, optional
        Small constant for numerical stability.

    Returns
    -------
    float
        Effective prototype count in [1, K].
    """
    pi_bar = cluster_weights.mean(dim=0)          # [K]
    h = -(pi_bar * torch.log(pi_bar + eps)).sum()
    return float(torch.exp(h).item())
