"""MoE-ViT: MoE++ (zero/copy/constant experts) on a pretrained DINOv3 ViT stem.

Replaces the MoR recursive core with a plain stacked-transformer core whose
FFNs are MoE++ layers (arXiv:2410.07348):

    dense MLP (d_ff=1536)  ->  4 FFN experts (d_ff=384) + 1 zero + 1 copy
                                + 2 constant experts, top_k=2 per token

Routing is per patch token, so easy image regions land on zero/copy/constant
specialists (near-zero compute) while hard regions route to real FFN experts -
adaptive compute at the finest granularity, with no recursion machinery.

The frozen stem and the study-level pooling/head mirror MoR-ViT so the two
models share the same ``features()`` / ``forward()`` contract and can be
compared on identical benchmarks.
"""

from __future__ import annotations

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

try:
    import timm

    HAS_TIMM = True
except ImportError:  # pragma: no cover - optional dependency
    HAS_TIMM = False


class SwiGLU(nn.Module):
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class ConstantExpert(nn.Module):
    """MoE++ constant expert: alpha1*x + alpha2*v (negligible params)."""

    def __init__(self, d_model: int):
        super().__init__()
        self.v = nn.Parameter(torch.zeros(d_model))
        self.wc = nn.Linear(d_model, 2, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        alpha = F.softmax(self.wc(x), dim=-1)
        return (alpha[..., 0:1] * x + alpha[..., 1:2] * self.v.to(dtype=x.dtype)).to(dtype=x.dtype)


class MoEFFN(nn.Module):
    """Sparse MoE FFN with MoE++ zero / copy / constant experts.

    Index layout: [FFN | zero | copy | constant].  Routing is softmax top-k
    with renormalisation; only tokens selected for a real FFN expert are
    computed (``index_select``), so zero-compute routing is real skipping.
    Returns the load-balance aux loss and pre-softmax gate logits.
    """

    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_ffn: int = 4,
        top_k: int = 2,
        n_zero: int = 1,
        n_copy: int = 1,
        n_const: int = 2,
        tau: float = 0.75,
        gating_residual: bool = True,
        gate_ctx: bool = False,
    ):
        super().__init__()
        self.num_ffn = num_ffn
        self.top_k = min(top_k, num_ffn + n_zero + n_copy + n_const)
        self.n_zero = max(0, n_zero)
        self.n_copy = max(0, n_copy)
        self.n_const = max(0, n_const)
        self.n_zc = self.n_zero + self.n_copy + self.n_const
        self.total_experts = num_ffn + self.n_zc
        self.last_counts: torch.Tensor | None = None
        self.tau = float(tau)
        self.gate_ctx = gate_ctx

        self.copy_start = num_ffn + self.n_zero
        self.const_start = self.copy_start + self.n_copy

        self.router = nn.Linear(d_model, self.total_experts, bias=False)
        self.experts = nn.ModuleList([SwiGLU(d_model, d_ff) for _ in range(num_ffn)])
        self.const_experts = nn.ModuleList([ConstantExpert(d_model) for _ in range(self.n_const)])

        self.gating_residual = None
        if gating_residual:
            self.gating_residual = nn.Linear(self.total_experts, self.total_experts, bias=False)
            nn.init.zeros_(self.gating_residual.weight)

        self.ctx_proj = None
        if gate_ctx:
            self.ctx_proj = nn.Linear(d_model, self.total_experts, bias=False)
            nn.init.zeros_(self.ctx_proj.weight)

        eta = torch.ones(self.total_experts)
        if self.n_zc > 0:
            eta[num_ffn:] = self.tau
        self.register_buffer("eta", eta, persistent=False)

    def forward(
        self,
        x: torch.Tensor,
        prev_gate: torch.Tensor | None = None,
        ctx: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Returns (output, aux_loss, gate_logits, gate_weights)."""
        B, T, D = x.shape
        x_2d = x.reshape(-1, D)
        N = x_2d.size(0)

        gate_logits = self.router(x_2d)
        if self.ctx_proj is not None and ctx is not None:
            gate_logits = gate_logits + self.ctx_proj(ctx).repeat_interleave(T, dim=0)
        if self.gating_residual is not None and prev_gate is not None:
            if prev_gate.dim() == 3:
                prev_gate = prev_gate.reshape(-1, prev_gate.size(-1))
            if prev_gate.size(0) == N and prev_gate.size(-1) == self.total_experts:
                gate_logits = gate_logits + self.gating_residual(prev_gate)

        gate_weights = F.softmax(gate_logits, dim=-1)
        top_weights, top_indices = gate_weights.topk(self.top_k, dim=-1)
        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True).clamp(min=1e-9)
        route_weights = top_weights.to(dtype=x_2d.dtype)

        flat_expert = top_indices.reshape(-1)
        flat_token = (
            torch.arange(N, device=x.device)
            .unsqueeze(1)
            .expand(-1, self.top_k)
            .reshape(-1)
        )
        flat_weight = route_weights.reshape(-1, 1)

        order = torch.argsort(flat_expert, stable=True)
        sorted_tokens = flat_token[order]
        sorted_weights = flat_weight[order]
        sorted_x = x_2d.index_select(0, sorted_tokens)
        counts = torch.bincount(
            flat_expert, minlength=self.total_experts
        ).to(torch.int64)
        offsets = torch.cat(
            [torch.zeros(1, dtype=torch.int64, device=x.device),
             torch.cumsum(counts, 0)[:-1]]
        )

        output = torch.zeros_like(x_2d)

        for expert_idx in range(self.num_ffn):
            lo = int(offsets[expert_idx])
            hi = lo + int(counts[expert_idx])
            if lo == hi:
                continue
            tids = sorted_tokens[lo:hi]
            w = sorted_weights[lo:hi]
            output.index_add_(
                0, tids,
                (self.experts[expert_idx](sorted_x[lo:hi]) * w).to(dtype=output.dtype),
            )

        if self.n_zc > 0:
            for i in range(self.n_copy):
                expert_idx = self.copy_start + i
                lo = int(offsets[expert_idx])
                hi = lo + int(counts[expert_idx])
                if lo == hi:
                    continue
                tids = sorted_tokens[lo:hi]
                w = sorted_weights[lo:hi]
                output.index_add_(
                    0, tids, (sorted_x[lo:hi] * w).to(dtype=output.dtype)
                )

            for i in range(self.n_const):
                expert_idx = self.const_start + i
                lo = int(offsets[expert_idx])
                hi = lo + int(counts[expert_idx])
                if lo == hi:
                    continue
                tids = sorted_tokens[lo:hi]
                w = sorted_weights[lo:hi]
                output.index_add_(
                    0, tids,
                    (self.const_experts[i](sorted_x[lo:hi]) * w).to(dtype=output.dtype),
                )

        counts = torch.zeros(self.total_experts, device=x.device, dtype=gate_weights.dtype)
        counts.scatter_add_(0, flat_expert, torch.ones_like(flat_expert, dtype=gate_weights.dtype))
        self.last_counts = counts.detach()
        fractions = counts / counts.sum().clamp(min=1e-12)
        avg_prob = gate_weights.mean(dim=0)
        if self.n_zc > 0:
            aux_loss = self.total_experts * (self.eta.to(dtype=gate_weights.dtype) * fractions * avg_prob).sum()
        else:
            aux_loss = self.total_experts * (fractions * avg_prob).sum()

        return (output.view(B, T, D), aux_loss, gate_logits,
                gate_weights.view(B, T, self.total_experts))


class MoEAttentionBlock(nn.Module):
    """Pre-norm attention + MoE++ FFN.  Stores the FFN gate maps for the
    spatial-smoothness (TV) loss."""

    def __init__(
        self,
        dim: int,
        n_heads: int = 6,
        d_ff: int = 384,
        num_ffn: int = 4,
        top_k: int = 2,
        n_zero: int = 1,
        n_copy: int = 1,
        n_const: int = 2,
        tau: float = 0.75,
        gating_residual: bool = True,
        gate_ctx: bool = False,
    ):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.q = nn.Linear(dim, dim)
        self.k = nn.Linear(dim, dim)
        self.v = nn.Linear(dim, dim)
        self.proj = nn.Linear(dim, dim)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.moe = MoEFFN(
            dim, d_ff, num_ffn=num_ffn, top_k=top_k,
            n_zero=n_zero, n_copy=n_copy, n_const=n_const,
            tau=tau, gating_residual=gating_residual, gate_ctx=gate_ctx,
        )

    def forward(
        self, x: torch.Tensor, prev_gate: torch.Tensor | None = None,
        ctx: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        S, T, D = x.shape
        H = self.n_heads
        xn = self.norm1(x)
        q = self.q(xn).reshape(S, T, H, self.head_dim).transpose(1, 2)
        k = self.k(xn).reshape(S, T, H, self.head_dim).transpose(1, 2)
        v = self.v(xn).reshape(S, T, H, self.head_dim).transpose(1, 2)
        attn = (q @ k.transpose(-1, -2) / (self.head_dim ** 0.5)).softmax(dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(S, T, D)
        x = x + self.proj(out)

        y, aux, gate_logits, gate_weights = self.moe(
            self.norm2(x), prev_gate=prev_gate, ctx=ctx)
        x = x + y
        return x, aux, gate_logits, gate_weights


class SliceAttentionPool(nn.Module):
    """Weight slices by learned relevance before aggregating into a study vector."""

    def __init__(self, dim: int):
        super().__init__()
        self.query = nn.Parameter(torch.randn(dim))
        self.scale = dim ** -0.5

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = F.softmax(x @ self.query * self.scale, dim=0)
        return (w.unsqueeze(1) * x).sum(0)


class PerTargetPool(nn.Module):
    """Per-finding attention over slices: each of n_classes findings owns a
    query that weights the slices most relevant to it -> [n_classes, D]."""

    def __init__(self, dim: int, n_classes: int):
        super().__init__()
        self.queries = nn.Parameter(torch.randn(n_classes, dim) * dim ** -0.5)
        self.scale = dim ** -0.5

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [S, D] slice vectors -> [n_classes, S] attention weights
        w = F.softmax(x @ self.queries.T * self.scale, dim=0)
        return (x.unsqueeze(0) * w.T.unsqueeze(-1)).sum(1)


class PerTargetHead(nn.Module):
    """Shared MLP then per-target projection (12 findings, each its own weight)."""

    def __init__(self, dim: int, n_classes: int, hidden: int = 256):
        super().__init__()
        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU())
        self.w = nn.Parameter(torch.randn(n_classes, hidden) * hidden ** -0.5)
        self.b = nn.Parameter(torch.zeros(n_classes))

    def forward(self, ctx: torch.Tensor) -> torch.Tensor:
        # ctx: [n_classes, D] -> logits [n_classes]
        h = self.mlp(ctx)
        return (h * self.w).sum(-1) + self.b


class MoEViT(nn.Module):
    """Pretrained DINOv3 stem + stacked MoE++ core with per-token routing."""

    def __init__(
        self,
        *,
        stem_name: str = "vit_small_patch16_dinov3",
        stem: nn.Module | None = None,
        pretrained: bool = True,
        freeze_stem: bool = True,
        core_blocks: int = 2,
        n_heads: int = 6,
        d_ff: int = 384,
        num_ffn: int = 4,
        top_k: int = 2,
        n_zero: int = 1,
        n_copy: int = 1,
        n_const: int = 2,
        tau: float = 0.75,
        gating_residual: bool = True,
        n_classes: int = 12,
        gate_ctx: bool = False,
        per_target: bool = False,
    ):
        super().__init__()
        if stem is None:
            if not HAS_TIMM:
                raise ImportError("timm is required for a pretrained stem")
            stem = timm.create_model(stem_name, pretrained=pretrained, num_classes=0)
        self.num_ffn = num_ffn
        self.top_k = top_k
        self.gate_ctx = gate_ctx
        self.per_target = per_target

        self.stem = stem
        stem_dim = self.stem.embed_dim
        self.core_dim = 384
        self.input_proj = nn.Linear(stem_dim, self.core_dim) if stem_dim != 384 else None
        dim = self.core_dim
        self.n_prefix = int(getattr(self.stem, "num_prefix_tokens", 1))

        if freeze_stem:
            for p in self.stem.parameters():
                p.requires_grad_(False)

        self._ffn_maps: list[torch.Tensor] = []

        self.blocks = nn.ModuleList(
            [
                MoEAttentionBlock(
                    dim, n_heads=n_heads, d_ff=d_ff, num_ffn=num_ffn, top_k=top_k,
                    n_zero=n_zero, n_copy=n_copy, n_const=n_const,
                    tau=tau, gating_residual=gating_residual, gate_ctx=gate_ctx,
                )
                for _ in range(core_blocks)
            ]
        )

        self.exit_norm = nn.LayerNorm(dim)
        if per_target:
            self.pool = PerTargetPool(dim, n_classes)
            self.head = PerTargetHead(dim, n_classes)
        else:
            self.pool = SliceAttentionPool(dim)
            self.head = nn.Sequential(
                nn.Linear(dim, dim * 2),
                nn.GELU(),
                nn.Linear(dim * 2, n_classes),
            )

    def moe_aux_loss(self) -> torch.Tensor:
        total = None
        for b in self.blocks:
            if not hasattr(b, "_aux"):
                continue
            total = b._aux if total is None else total + b._aux
        if total is None:
            return torch.tensor(0.0, device=next(self.parameters()).device)
        return total / len(self.blocks)

    def routing_tv_loss(self) -> torch.Tensor:
        """Total-variation penalty on the per-patch FFN-mass maps.

        Knee findings occupy contiguous regions; scattered routing is a bug
        signal.  The map is the probability mass routed to real FFN experts,
        reshaped to the patch grid.
        """
        if not self._ffn_maps:
            return torch.tensor(0.0, device=next(self.parameters()).device)
        total = None
        for p in self._ffn_maps:
            P = p.size(1)
            h = int(round(P ** 0.5))
            if h * h != P:
                continue
            g = p.reshape(p.size(0), h, h)
            tv = (g[:, 1:, :] - g[:, :-1, :]).abs().mean() + (
                g[:, :, 1:] - g[:, :, :-1]
            ).abs().mean()
            total = tv if total is None else total + tv
        if total is None:
            return torch.tensor(0.0, device=next(self.parameters()).device)
        return total / len(self._ffn_maps)

    def features(
        self, x: torch.Tensor, pool: str = "patchmean"
    ) -> torch.Tensor:
        """Per-slice features (no head): [S, D] from [S, 3, H, W] slices."""
        tokens = self.stem.forward_features(x)  # [S, T, D]
        return self.features_from_tokens(tokens, pool)

    def features_from_tokens(
        self, tokens: torch.Tensor, pool: str = "patchmean"
    ) -> torch.Tensor:
        """Features from cached stem tokens (feat cache path): [S, D]."""
        tokens = tokens.float()  # core internals are f32; bf16 cache upcasts exactly
        if self.input_proj is not None:
            with torch.autocast("cuda", enabled=False):
                tokens = self.input_proj(tokens).float()  # keep f32; autocast would make it bf16
        self._ffn_maps = []
        prev_gate = None
        ctx = tokens.mean(dim=1) if self.gate_ctx else None  # [B, D] per-slice context
        for b in self.blocks:
            tokens, aux, gate_logits, gate_weights = b(
                tokens, prev_gate=prev_gate, ctx=ctx)
            b._aux = aux
            self._ffn_maps.append(
                gate_weights[:, self.n_prefix:, : self.num_ffn].sum(dim=-1)
            )
            prev_gate = gate_logits
        tokens = self.exit_norm(tokens)
        if pool == "cls":
            return tokens[:, 0]
        return tokens[:, self.n_prefix:].mean(1)

    def forward(
        self, x: torch.Tensor
    ) -> tuple[torch.Tensor, list[float]]:
        """Forward one study: [S, 3, H, W] slices -> (logits [n_classes], aux_stats)."""
        slice_feats = self.features(x)  # [S, D]
        study = self.pool(slice_feats)  # [D] or [n_classes, D]
        logits = self.head(study)
        if logits.dim() == 2:
            logits = logits.squeeze(-1)
        return logits, [self.moe_aux_loss().item()]