"""panda encoder with two input variants (pca / marker), sub-center prototypes."""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import Optional, List

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function


# gradient reversal layer

class GradReverse(Function):
    @staticmethod
    def forward(ctx, x, lam):
        ctx.lam = lam
        return x.view_as(x)

    @staticmethod
    def backward(ctx, g):
        return -ctx.lam * g, None


def grad_reverse(x, lam):
    return GradReverse.apply(x, lam)


# encoder with variant + sub-centers

class PANDAEncoder(nn.Module):
    """trunk + projection head + sub-center prototypes.

    args:
      variant     : "pca" or "marker"
      n_pca       : 50
      n_markers   : m >= 0. required > 0 if variant == "marker".
      n_classes   : K
      n_sub       : sub-centers per class (default 3)
      d_hidden, d_repr, d_proj: trunk sizing
      n_datasets  : for the dataset adversary head
    """

    def __init__(
        self,
        variant: str = "pca",
        n_pca: int = 50,
        n_markers: int = 0,
        d_hidden: int = 512,
        d_repr: int = 256,
        d_proj: int = 128,
        n_classes: int = 10,
        n_sub: int = 3,
        n_datasets: int = 1,
        dropout: float = 0.2,
    ):
        super().__init__()
        assert variant in ("pca", "marker"), variant
        if variant == "marker":
            assert n_markers > 0, "PANDA-Marker requires n_markers>0"
        self.variant = variant
        self.n_pca = n_pca
        self.n_markers = n_markers if variant == "marker" else 0
        self.n_classes = n_classes
        self.n_sub = n_sub
        self.n_datasets = n_datasets

        input_dim = n_pca + self.n_markers
        self.input_dim = input_dim

        self.trunk = nn.Sequential(
            nn.Linear(input_dim, d_hidden), nn.LayerNorm(d_hidden), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(d_hidden, d_hidden), nn.LayerNorm(d_hidden), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(d_hidden, d_repr), nn.LayerNorm(d_repr), nn.GELU(),
        )
        self.projection = nn.Sequential(
            nn.Linear(d_repr, d_repr), nn.GELU(),
            nn.Linear(d_repr, d_proj),
        )
        self.classifier = nn.Sequential(nn.Linear(d_repr + 2, n_classes))
        self.dom_adv = nn.Sequential(nn.Linear(d_repr, 128), nn.ReLU(), nn.Linear(128, n_datasets))
        self.depth_adv = nn.Sequential(nn.Linear(d_repr, 64), nn.ReLU(), nn.Linear(64, 1))

        # sub-center prototypes (K, n_sub, d_proj), L2-normalised per sub-center
        self.register_buffer(
            "prototypes",
            F.normalize(torch.randn(n_classes, n_sub, d_proj), dim=-1),
        )
        # EMA momentum as a buffer so we can overwrite it in place
        self.register_buffer("proto_ema", torch.tensor(0.99))

    @torch.no_grad()
    def update_prototypes(self, z_norm: torch.Tensor, y: torch.Tensor):
        """ema update: assign each in-class cell to nearest sub-center, take the mean."""
        ema = float(self.proto_ema.item())
        for c in torch.unique(y):
            mask = y == c
            if not mask.any():
                continue
            zc = z_norm[mask]                              # (n_c, d_proj)
            protos_c = self.prototypes[c]                  # (n_sub, d_proj)
            sims = zc @ protos_c.T                         # (n_c, n_sub)
            assign = sims.argmax(dim=1)                    # each cell -> nearest sub-center
            for k in range(self.n_sub):
                m2 = assign == k
                if not m2.any():
                    continue
                new = F.normalize(zc[m2].mean(dim=0), dim=0)
                self.prototypes[c, k] = F.normalize(
                    ema * self.prototypes[c, k] + (1 - ema) * new, dim=0
                )

    @torch.no_grad()
    def max_sub_cos(self, z_norm: torch.Tensor) -> torch.Tensor:
        """(B, K) cos(z, best sub-center) per class."""
        B = z_norm.size(0); K, n_sub, D = self.prototypes.shape
        sims = torch.einsum("bd,ksd->bks", z_norm, self.prototypes)   # (B, K, n_sub)
        return sims.max(dim=2).values                                 # (B, K)

    def forward(
        self,
        x_pca: torch.Tensor,
        aux: torch.Tensor,
        x_markers: Optional[torch.Tensor] = None,
        lam_dann: float = 0.0,
    ) -> dict:
        if self.variant == "marker":
            assert x_markers is not None and x_markers.size(1) == self.n_markers
            x = torch.cat([x_pca, x_markers], dim=1)
        else:
            x = x_pca

        h = self.trunk(x)
        z_raw = self.projection(h)
        z = F.normalize(z_raw, dim=1)
        logits = self.classifier(torch.cat([h, aux], dim=1))
        h_rev = grad_reverse(h, lam_dann)
        return {
            "repr": h,
            "z": z,
            "logits": logits,
            "dom": self.dom_adv(h_rev),
            "depth": self.depth_adv(h_rev),
        }


# losses

def supcon_loss(z: torch.Tensor, y: torch.Tensor, temperature: float = 0.1) -> torch.Tensor:
    if z.size(0) < 2:
        return z.new_zeros(())
    sim = z @ z.T / temperature
    sim_max, _ = sim.max(dim=1, keepdim=True)
    sim = sim - sim_max.detach()
    logits_mask = torch.ones_like(sim) - torch.eye(z.size(0), device=z.device)
    exp_sim = torch.exp(sim) * logits_mask
    log_prob = sim - torch.log(exp_sim.sum(dim=1, keepdim=True) + 1e-12)
    labels_eq = (y.unsqueeze(0) == y.unsqueeze(1)).float() * logits_mask
    denom = labels_eq.sum(dim=1).clamp_min(1.0)
    per = -(labels_eq * log_prob).sum(dim=1) / denom
    per = per * (labels_eq.sum(dim=1) > 0).float()
    counts = torch.bincount(y, minlength=int(y.max().item()) + 1).float().clamp_min(1.0)
    w = 1.0 / counts.sqrt()
    return (per * w[y]).sum() / w[y].sum().clamp_min(1e-6)


def vicreg_loss(z: torch.Tensor, sim_weight: float = 0.0, var_weight: float = 25.0,
                cov_weight: float = 1.0) -> torch.Tensor:
    zc = z - z.mean(dim=0, keepdim=True)
    std = (zc.var(dim=0) + 1e-4).sqrt()
    var_loss = F.relu(1.0 - std).mean()
    N, D = zc.shape
    cov = (zc.T @ zc) / (N - 1)
    off = cov - torch.diag(torch.diagonal(cov))
    cov_loss = off.pow(2).sum() / D
    return var_weight * var_loss + cov_weight * cov_loss


def hsic_biased(x: torch.Tensor, y: torch.Tensor,
                sigma_x: float = 1.0, sigma_y: float = 1.0) -> torch.Tensor:
    Nx = x.size(0)
    if Nx < 2:
        return x.new_zeros(())
    K = torch.exp(-torch.cdist(x, x) ** 2 / (2 * sigma_x ** 2))
    L = torch.exp(-torch.cdist(y, y) ** 2 / (2 * sigma_y ** 2))
    H = torch.eye(Nx, device=x.device) - torch.ones(Nx, Nx, device=x.device) / Nx
    return (K @ H @ L @ H).trace() / (Nx - 1) ** 2


def subcenter_angular_infonce(
    z: torch.Tensor,                    # (B, d_proj) L2-normalised
    y: torch.Tensor,                    # (B,)
    prototypes: torch.Tensor,           # (K, n_sub, d_proj)
    margin: float = 0.15,               # angular margin in radians
    temperature: float = 0.07,
) -> torch.Tensor:
    """arcface-style angular-margin loss over sub-center prototypes."""
    B = z.size(0); K, n_sub, D = prototypes.shape
    sims = torch.einsum("bd,ksd->bks", z, prototypes)          # (B, K, n_sub)
    max_over_sub = sims.max(dim=2).values                       # (B, K)

    # target class cosine, bump by angular margin, put back
    target_cos = max_over_sub.gather(1, y.unsqueeze(1)).squeeze(1)   # (B,)
    target_cos = target_cos.clamp(-1 + 1e-7, 1 - 1e-7)
    theta = torch.acos(target_cos)
    target_new_cos = torch.cos(theta + margin)

    logits = max_over_sub.clone()
    logits.scatter_(1, y.unsqueeze(1), target_new_cos.unsqueeze(1))
    logits = logits / temperature
    return F.cross_entropy(logits, y)


def prototype_repulsion(prototypes: torch.Tensor, weight: float = 1.0) -> torch.Tensor:
    """penalise inter-class prototype cosine so eff-dim doesn't collapse."""
    K, n_sub, D = prototypes.shape
    centroids = F.normalize(prototypes.mean(dim=1), dim=1)      # (K, D)
    sim = centroids @ centroids.T                                # (K, K)
    off = sim - torch.diag(torch.diagonal(sim))
    return weight * off.pow(2).sum() / (K * (K - 1) + 1e-6)


def prototype_infonce_legacy(z, y, prototypes, temperature=0.07):
    """legacy single-prototype InfoNCE. kept for debugging + old checkpoints."""
    if prototypes.dim() == 3:
        prototypes = F.normalize(prototypes.mean(dim=1), dim=1)  # collapse sub-centers
    logits = z @ prototypes.T / temperature
    return F.cross_entropy(logits, y)