"""
SuperCluster: Learning Prediction Prototypes via Target-Guided Cross-Attention Clustering.

Core architecture:

  x  →  MLP encoder  →  z^(0)                            [N, d]
  z^(0..L)  →  L × CrossTransformerLayer                  [N, d]   (attends to centers C)
  z^(L)  →  attn_to_logits  →  softmax(·/τ)  →  π_i     [N, K]   routing weights

  Binary (num_classes=1):
    ℓ_i = π_i · s,   ŷ_i = σ(ℓ_i)       [N, 1]    s ∈ R^K

  Multi-class (num_classes=C):
    ℓ_i = π_i @ S,   ŷ_i = softmax(ℓ_i) [N, C]    S ∈ R^{K×C}

  Regression: same as binary, replace BCEWithLogitsLoss → MSELoss, no sigmoid at inference.

Key design: prototype centers C (routing geometry) and prototype scores s/S (prediction)
are separate parameters, decoupling diversity regularization from prediction confidence.
"""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F


class SuperCluster(nn.Module):
    """Supervised prototype-routing model for tabular classification and regression.

    Learns K prediction prototypes jointly with a supervised objective.
    Examples are routed to prototypes via cross-attention; the prediction is a
    convex combination of prototype-level predictions (scalars for binary/regression,
    class-logit vectors for multi-class).

    Parameters
    ----------
    input_dim : int
        Raw feature dimension.
    embedding_dim : int
        Latent embedding dimension d. Typical values: 64–256.
    num_clusters : int
        Number of prototypes K. Set larger than expected; surplus prototypes
        will collapse to zero routing mass during training.
    num_classes : int, optional
        1 for binary classification or regression (default).
        C > 1 for C-class multi-class classification.
    encoder_layers : int, optional
        Depth of the MLP encoder (default 4).
    encoder_hidden_size : int, optional
        Hidden size for the MLP encoder (default 256).
    d_ff : int, optional
        Feed-forward hidden size inside each CrossTransformerLayer.
        Defaults to 4 * embedding_dim (standard transformer convention).
    num_attn_heads : int, optional
        Number of attention heads (default 8). Must divide embedding_dim.
    num_cross_attn_layers : int, optional
        Number of cross-attention + FFN layers L (default 2).
    dropout : float, optional
        Dropout rate (default 0.1).
    lambda_clust : float, optional
        Weight for the internal clustering regularizer (default 0.0).
        In practice the center-diversity loss is applied externally in training
        via :func:`supercluster.losses.center_diversity_loss`.

    Examples
    --------
    Binary classification::

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

        model = SuperCluster(input_dim=10, embedding_dim=64, num_clusters=8)
        x = torch.randn(32, 10)
        logits, cluster_weights = model(x)   # logits: [32, 1], weights: [32, 8]
        loss = nn.BCEWithLogitsLoss()(logits, y)

    7-class classification::

        model = SuperCluster(input_dim=54, embedding_dim=128,
                             num_clusters=8, num_classes=7)
        logits, cluster_weights = model(x)   # logits: [32, 7]
        loss = nn.CrossEntropyLoss()(logits, y_long)

    Regression::

        model = SuperCluster(input_dim=10, embedding_dim=64, num_clusters=5)
        preds, cluster_weights = model(x)    # preds: [32, 1], no sigmoid
        loss = nn.MSELoss()(preds, y)

    Reading prototype predictions::

        # Binary / regression: σ(s_k) or s_k directly
        for k in range(model.num_clusters):
            print(f"Prototype {k}: s_k = {model.prototype_scores[k].item():.3f}")

        # Multi-class: softmax(S_k) is the predicted class distribution
        import torch.nn.functional as F
        for k in range(model.num_clusters):
            probs = F.softmax(model.prototype_scores[k], dim=0)
            print(f"Prototype {k}: {probs.tolist()}")
    """

    def __init__(
        self,
        input_dim: int,
        embedding_dim: int,
        num_clusters: int,
        num_classes: int = 1,
        encoder_layers: int = 4,
        encoder_hidden_size: int = 256,
        d_ff: int | None = None,
        num_attn_heads: int = 8,
        num_cross_attn_layers: int = 2,
        dropout: float = 0.1,
        lambda_clust: float = 0.0,
    ) -> None:
        super().__init__()
        self.input_dim = input_dim
        self.embedding_dim = embedding_dim
        self.num_clusters = num_clusters
        self.num_classes = num_classes
        self.lambda_clust = lambda_clust

        if d_ff is None:
            d_ff = 4 * embedding_dim

        # Learnable routing temperature (initialised at 1.0; converges to ~0.5–0.8)
        self.temperature = nn.Parameter(torch.tensor(1.0))

        # Prototype center embeddings C ∈ R^{K × d} — routing geometry only
        self.centers = nn.Parameter(torch.randn(num_clusters, embedding_dim))

        self.encoder = MLPBlock(
            input_dim, embedding_dim,
            num_layers=encoder_layers,
            hidden_size=encoder_hidden_size,
            dropout=dropout,
        )

        self.cross_attn_layers = nn.ModuleList([
            CrossTransformerLayer(embedding_dim, num_attn_heads, d_ff, dropout)
            for _ in range(num_cross_attn_layers)
        ])

        self.attn_to_logits = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim),
            nn.ReLU(),
            nn.Linear(embedding_dim, num_clusters),
        )

        # Prototype prediction scores — separate from routing geometry.
        # Binary / regression: s ∈ R^K (one scalar per prototype)
        # Multi-class:         S ∈ R^{K × C} (one logit vector per prototype)
        if num_classes == 1:
            self.prototype_scores = nn.Parameter(torch.zeros(num_clusters))
        else:
            self.prototype_scores = nn.Parameter(torch.zeros(num_clusters, num_classes))

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Forward pass.

        Parameters
        ----------
        x : torch.Tensor, shape [N, input_dim]
            Batch of input features.

        Returns
        -------
        logits : torch.Tensor
            Binary/regression → [N, 1] (ℓ_i = π_i · s; pass to BCEWithLogitsLoss
            or MSELoss).
            Multi-class → [N, C] (ℓ_i = π_i @ S; pass to CrossEntropyLoss).
        cluster_weights : torch.Tensor, shape [N, K]
            Soft cluster assignment probabilities π_i ∈ Δ^{K-1}.
        """
        queries = F.normalize(self.encoder(x), dim=-1)              # [N, d]
        centers = F.normalize(self.centers, dim=-1)                  # [K, d]

        z = queries
        for layer in self.cross_attn_layers:
            z = layer(z, centers)                                    # [N, d]

        logits_k = self.attn_to_logits(z)                           # [N, K]
        cluster_weights = F.softmax(logits_k / self.temperature, dim=-1)  # [N, K]

        if self.num_classes == 1:
            # Binary / regression: ℓ_i = Σ_k π_ik · s_k  →  [N, 1]
            pred_logits = (cluster_weights * self.prototype_scores).sum(
                dim=-1, keepdim=True
            )
        else:
            # Multi-class: ℓ_i = π_i @ S  →  [N, C]
            pred_logits = cluster_weights @ self.prototype_scores

        return pred_logits, cluster_weights

    def active_prototypes(
        self,
        cluster_weights: torch.Tensor,
        threshold: float = 0.01,
    ) -> list[int]:
        """Return indices of prototypes that receive >= threshold mean routing mass.

        Parameters
        ----------
        cluster_weights : torch.Tensor, shape [N, K]
            Soft assignment weights from a forward pass.
        threshold : float, optional
            Minimum mean routing mass to be considered active (default 0.01).

        Returns
        -------
        list[int]
            Sorted list of active prototype indices.
        """
        mean_dist = cluster_weights.mean(dim=0)
        return sorted((mean_dist >= threshold).nonzero(as_tuple=True)[0].tolist())

    def clustering_loss(
        self, cluster_weights: torch.Tensor, eps: float = 1e-8
    ) -> torch.Tensor:
        """Entropy-based clustering regulariser (optional; center_diversity_loss preferred).

        Encourages sharp per-example assignments and balanced usage across prototypes:
        L = E[H(π_i)] - H(mean_i π_i).

        Parameters
        ----------
        cluster_weights : torch.Tensor, shape [N, K]
        eps : float, optional
            Small constant for numerical stability.

        Returns
        -------
        torch.Tensor
            Scalar regularisation loss.
        """
        per_sample_entropy = (
            -(cluster_weights * torch.log(cluster_weights + eps)).sum(dim=1).mean()
        )
        mean_dist = cluster_weights.mean(dim=0)
        batch_entropy = -(mean_dist * torch.log(mean_dist + eps)).sum()
        return per_sample_entropy - batch_entropy


class CrossTransformerLayer(nn.Module):
    """One cross-attention + feed-forward block.

    Queries come from example embeddings; keys/values come from prototype centers.

    Parameters
    ----------
    d_model : int
        Model dimension.
    n_heads : int
        Number of attention heads. Must divide d_model.
    d_ff : int
        Feed-forward hidden size (typically 4 * d_model).
    dropout : float
        Dropout rate.
    """

    def __init__(
        self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1
    ) -> None:
        super().__init__()
        self.cross_attn = CrossAttention(d_model, n_heads, dropout=dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, X: torch.Tensor, C: torch.Tensor) -> torch.Tensor:
        """
        Parameters
        ----------
        X : torch.Tensor, shape [N, d_model]
            Query embeddings (examples).
        C : torch.Tensor, shape [K, d_model]
            Key/value embeddings (prototype centers).

        Returns
        -------
        torch.Tensor, shape [N, d_model]
        """
        attn_out = self.cross_attn(X, C)
        X = self.norm1(X + self.dropout1(attn_out))
        return self.norm2(X + self.dropout2(self.ff(X)))


class CrossAttention(nn.Module):
    """Scaled dot-product multi-head cross-attention.

    Parameters
    ----------
    d_model : int
    n_heads : int
    dropout : float
    """

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1) -> None:
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.Wq = nn.Linear(d_model, d_model)
        self.Wk = nn.Linear(d_model, d_model)
        self.Wv = nn.Linear(d_model, d_model)
        self.out = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, X: torch.Tensor, C: torch.Tensor) -> torch.Tensor:
        N, M = X.size(0), C.size(0)
        Q = self.Wq(X).view(N, self.n_heads, self.d_k)
        K = self.Wk(C).view(M, self.n_heads, self.d_k)
        V = self.Wv(C).view(M, self.n_heads, self.d_k)
        scores = torch.einsum("nhd,mhd->nhm", Q, K) / (self.d_k**0.5)
        alpha = self.dropout(F.softmax(scores, dim=-1))              # [N, H, M]
        out = torch.einsum("nhm,mhd->nhd", alpha, V)
        return self.out(out.contiguous().view(N, self.d_model))


class MLPBlock(nn.Module):
    """Residual MLP with LayerNorm and ReLU activations.

    Parameters
    ----------
    input_size : int
    output_size : int
    num_layers : int
        Total depth including input and output projections.
    hidden_size : int
    dropout : float
    """

    def __init__(
        self,
        input_size: int,
        output_size: int,
        num_layers: int = 4,
        hidden_size: int = 64,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        layers: list[nn.Module] = [
            nn.Linear(input_size, hidden_size),
            nn.LayerNorm(hidden_size),
            nn.ReLU(),
            nn.Dropout(dropout),
        ]
        for _ in range(num_layers - 2):
            layers += [
                nn.Linear(hidden_size, hidden_size),
                nn.LayerNorm(hidden_size),
                nn.ReLU(),
                nn.Dropout(dropout),
            ]
        layers.append(nn.Linear(hidden_size, output_size))
        self.model = nn.Sequential(*layers)
        self.skip = (
            nn.Linear(input_size, output_size)
            if input_size != output_size
            else nn.Identity()
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.skip(x) + self.model(x)


# Backward-compatibility alias for code using the old name
SupervisedClusteringModelV2 = SuperCluster
