src/models/plant_vit.py

"""PlantViT - thin adapter over RSNA MoR/MoE for single-image plant ID.
Reuses moe_vit.MoEViT and mor_vit.MoRViT frozen DINOv3 stem + routing.
Single image = single study with S=1 slice, so pool reduces to patchmean.
Adds 512-D embedder head per architecture.md §4.

Usage:
  from src.models.plant_vit import PlantViT
  model = PlantViT(stem_name="vit_base_patch16_dinov3", n_classes=500, use_moe=True, embed_dim=512)
  logits, embed, aux = model(x)  # x [B,3,336,336] -> logits [B,500], embed [B,512] L2
"""

from __future__ import annotations

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

from .moe_vit import MoEViT
from .mor_vit import MoRViT


class EmbedHead(nn.Module):
    """Plant-specific projection: 768/384 -> 512 L2 per architecture.md §4."""
    def __init__(self, in_dim: int, embed_dim: int = 512):
        super().__init__()
        self.ln1 = nn.LayerNorm(in_dim)
        self.proj = nn.Linear(in_dim, embed_dim, bias=False)
        self.ln2 = nn.LayerNorm(embed_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x [B, D]
        x = self.ln1(x)
        x = self.proj(x)
        x = self.ln2(x)
        return F.normalize(x, dim=-1)


class PlantViT(nn.Module):
    """Wrapper that exposes (logits, embed, aux_losses) with interchangeable stem.

    Args:
        stem_name: timm name, e.g. vit_small_patch16_dinov3, vit_base_patch16_dinov3, vit_large_patch16_dinov3, convnext_tiny
        n_classes: 500 for WA top500
        embed_dim: 512 (test 256/128 later)
        use_moe: True -> MoE++ (MoEViT), False -> MoR (MoRViT), None -> dense
        freeze_stem: keep DINOv3 frozen for CE baseline
        All MoE/MoR kwargs (num_ffn, top_k, d_ff, n_recursions, capacities etc.) passed through.
    """
    def __init__(
        self,
        *,
        stem_name: str = "vit_base_patch16_dinov3",
        n_classes: int = 500,
        embed_dim: int = 512,
        use_moe: bool | None = True,
        freeze_stem: bool = True,
        # MoE knobs
        num_ffn: int = 4,
        top_k: int = 2,
        d_ff: int = 384,
        # MoR knobs
        n_recursions: int = 3,
        capacities: tuple[float, ...] = (1.0, 0.667, 0.333),
        core_blocks: int = 2,
        router_hidden: int = 128,
        pretrained: bool = True,
        **kwargs,
    ):
        super().__init__()
        self.use_moe = use_moe
        # build core - keep RSNA defaults, just swap n_classes and stem
        if use_moe is True:
            self.core = MoEViT(
                stem_name=stem_name, pretrained=pretrained, freeze_stem=freeze_stem,
                core_blocks=core_blocks, d_ff=d_ff, num_ffn=num_ffn, top_k=top_k,
                n_classes=n_classes, **{k: v for k, v in kwargs.items() if k in ["gate_ctx","per_target","tau","n_zero","n_copy","n_const","gating_residual"]},
            )
            feat_dim = self.core.core_dim  # 384
        elif use_moe is False:
            self.core = MoRViT(
                stem_name=stem_name, pretrained=pretrained, freeze_stem=freeze_stem,
                n_recursions=n_recursions, capacities=capacities, core_blocks=core_blocks,
                router_hidden=router_hidden, n_classes=n_classes,
            )
            feat_dim = self.core.core_dim
        else:
            # dense baseline - MoR dense=True
            self.core = MoRViT(
                stem_name=stem_name, pretrained=pretrained, freeze_stem=freeze_stem,
                n_recursions=n_recursions, capacities=capacities, core_blocks=core_blocks,
                router_hidden=router_hidden, n_classes=n_classes, dense=True,
            )
            feat_dim = self.core.core_dim

        # embedder on pooled slice feat (before head pool, use same dim)
        # For MoE/MoR, slice feat dim = core.core_dim (384) ; for base ViT-B 768 -> still 384 via input_proj
        # If stem is vit_base (768) we still route through input_proj -> 384, so embed in 384
        self.embed = EmbedHead(feat_dim, embed_dim)
        self.n_classes = n_classes
        self.embed_dim = embed_dim
        self.stem_name = stem_name

    def forward(self, x: torch.Tensor, step: int = 0):
        """x [B,3,H,W] single image per plant -> logits [B,C], embed [B,512], aux dict"""
        # RSNA core expects [S,3,H,W] where S is slices per study; for plant S=B single slice per sample
        # Use core.features single-batch path
        feats = self.core.features(x)  # [B, D]  (MoEViT/MoRViT handle bfloat16 internally)
        logits = self.core.head(self.core.pool(feats.unsqueeze(1).repeat(1,1,1)) if False else feats)  # bypass buggy repeat - just pool slice feats
        # Actually MoEViT/MoRViT head expects pooled study vector [D]; we have [B,D] already pooled per image
        # Re-derive logits via head directly for single-slice case:
        # MoEViT/MoRViT: features -> [S,D] -> pool -> [D] -> head -> [C]
        # For B independent images, features is [B,D] already per-image slice feat, so head per row:
        # Use core.head on each row individually via batched linear - head is Sequential Linear/GELU/Linear
        # For per_target head, need [B, n_classes, D] - not used for plant (per_target=False by default)
        # Simpler: call core.head on feats directly (feats is [B,D] -> head expects [D] but Linear broadcasts on last dim)
        # Do batched:
        if isinstance(self.core.head, nn.Sequential):
            logits = self.core.head(feats)  # [B,C] broadcasts
        else:
            # PerTarget not used
            logits = self.core.head(feats)

        embed = self.embed(feats)
        aux = {}
        # surface aux losses if present
        try:
            aux["moe_aux"] = float(self.core.moe_aux_loss().item()) if hasattr(self.core, "moe_aux_loss") else 0.0
        except:
            aux["moe_aux"] = 0.0
        try:
            aux["tv"] = float(self.core.routing_tv_loss().item()) if hasattr(self.core, "routing_tv_loss") else float(self.core.routing_smoothness_loss().item() if hasattr(self.core, "routing_smoothness_loss") else 0.0)
        except:
            aux["tv"] = 0.0
        return logits, embed, aux

    def features(self, x: torch.Tensor) -> torch.Tensor:
        return self.core.features(x)