"""
Sub-Center ArcFace with K Sub-Centers per Class.
Designed for fine-grained botanical vision to capture multi-modal phenology
(e.g., Flower vs. Leaf/Foliage vs. Fruit/Seed Pods) without intra-class noise.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class SubCenterArcFace(nn.Module):
    """
    Sub-Center ArcFace for 512-D L2 Embeddings.
    Maintains K sub-centers per class (default K=3).
    """
    def __init__(self, embed_dim: int, n_classes: int, k_subcenters: int = 3, m: float = 0.5, s: float = 30.0, easy_margin: bool = False):
        super().__init__()
        self.embed_dim = embed_dim
        self.n_classes = n_classes
        self.k_subcenters = k_subcenters
        self.m = m
        self.s = s
        self.easy_margin = easy_margin

        # Weight matrix: [n_classes * k_subcenters, embed_dim]
        self.W = nn.Parameter(torch.randn(n_classes * k_subcenters, embed_dim) * 0.01)
        nn.init.xavier_uniform_(self.W)

        # Precomputed angular constants
        self.cos_m = math.cos(m)
        self.sin_m = math.sin(m)
        self.th = math.cos(math.pi - m)
        self.mm = math.sin(math.pi - m) * m

    def forward(self, emb: torch.Tensor, y: torch.Tensor = None) -> torch.Tensor:
        """
        emb: [B, D] L2 normalized embedding
        y: [B] Ground-truth class labels (0 .. n_classes-1)
        Returns: [B, n_classes] Logits scaled by s
        """
        # Normalize weights and embeddings
        W_norm = F.normalize(self.W, dim=1) # [C * K, D]
        emb_norm = F.normalize(emb, dim=1)   # [B, D]

        # Compute cosine similarity across all sub-centers: [B, C * K]
        cos_all = F.linear(emb_norm, W_norm)
        B = emb.size(0)

        # Reshape to [B, C, K] and select the maximum cosine per class
        cos_sub = cos_all.view(B, self.n_classes, self.k_subcenters)
        cos_max, _ = cos_sub.max(dim=-1) # [B, C]
        cos = cos_max.clamp(-1.0 + 1e-7, 1.0 - 1e-7)

        # Inference mode: return scaled max cosine logits
        if y is None:
            return cos * self.s

        # Training mode: Apply additive angular margin to ground-truth class
        sin = torch.sqrt((1.0 - cos ** 2).clamp(0.0, 1.0))
        cos_m = cos * self.cos_m - sin * self.sin_m

        if self.easy_margin:
            cos_m = torch.where(cos > 0, cos_m, cos)
        else:
            cos_m = torch.where(cos > self.th, cos_m, cos - self.mm)

        one_hot = F.one_hot(y, num_classes=self.n_classes).bool()
        logits = torch.where(one_hot, cos_m, cos)
        return logits * self.s

    def diversity_loss(self) -> torch.Tensor:
        """
        Penalise sub-center collapse (only winning K=1 gets gradients via max-pool).
        Computes mean absolute pairwise cosine between sub-centers within each class;
        encourages K=3 centers to stay diverse (e.g., flower/leaf/fruit).
        """
        if self.k_subcenters <= 1:
            return torch.tensor(0.0, device=self.W.device, dtype=self.W.dtype)
        # Reshape to [C, K, D] and L2-normalize per sub-center
        W_ck = F.normalize(self.W.view(self.n_classes, self.k_subcenters, self.embed_dim), dim=-1)
        # Pairwise cosine within each class: [C, K, K]
        cos = torch.einsum("ckd,cjd->ckj", W_ck, W_ck)
        # Mask diagonal (self-similarity =1)
        mask = torch.eye(self.k_subcenters, device=cos.device, dtype=torch.bool).unsqueeze(0).expand(
            self.n_classes, -1, -1
        )
        cos_off = cos[~mask]
        if cos_off.numel() == 0:
            return torch.tensor(0.0, device=cos.device, dtype=cos.dtype)
        # Want low similarity → penalise absolute correlation; hinge variant: relu(cos - 0.2)
        # Use abs mean for symmetry (both positive collapse and opposite collapse are diverse enough?)
        # but positive collapse is the failure mode, so use mean of positive part.
        # Here we use abs mean scaled: keeps gradient if centers align or anti-align strongly.
        return cos_off.abs().mean()