"""MRI-CORE stem: MRI-pretrained ViT-B feature extractor.

Loads the MRI-CORE checkpoint (DINOv2-style ViT-B, SAM-init, pretrained on
6.1M MRI slices) and exposes a timm-compatible surface so the existing
MoR/MoE cores can use it as their frozen stem:

    stem.embed_dim          -> 768
    stem.num_prefix_tokens  -> 1  (CLS only, no registers)
    stem.pos_embed          -> [1, 197, 768]
    stem.forward_features(x)-> [S, 197, 768]  (frozen, eval mode)

Key layout in the checkpoint:  teacher.backbone.{cls_token, pos_embed,
patch_embed.proj, blocks.<g>.<i>, norm}  with blocks nested 4 x 3 = 12.
"""
from __future__ import annotations

import os

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

_CKPT = "MRI_CORE_vitb.pth"


class Block(nn.Module):
    def __init__(self, dim: int = 768, n_heads: int = 12, mlp_ratio: float = 4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, int(dim * mlp_ratio)),
            nn.GELU(),
            nn.Linear(int(dim * mlp_ratio), dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.mlp(self.norm2(x))
        return x


class MRICoreStem(nn.Module):
    """Frozen MRI-CORE ViT-B feature extractor (embed_dim=768, 197 tokens)."""

    def __init__(self, ckpt_path: str | None = None, freeze: bool = True):
        super().__init__()
        if ckpt_path is None:
            for base in (
                r"G:\RSNA-Knee\cache",
                os.path.join(os.path.dirname(__file__), "..", "..", "weights"),
            ):
                cand = os.path.join(base, _CKPT)
                if os.path.isfile(cand):
                    ckpt_path = cand
                    break
        if ckpt_path is None or not os.path.isfile(ckpt_path):
            raise FileNotFoundError(f"MRI-CORE checkpoint not found ({_CKPT})")
        self.embed_dim = 768
        self.num_prefix_tokens = 1
        self.patch_embed = nn.Conv2d(3, 768, kernel_size=16, stride=16, bias=True)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, 768))
        self.pos_embed = nn.Parameter(torch.zeros(1, 197, 768))
        self.mask_token = nn.Parameter(torch.zeros(1, 768))
        self.blocks = nn.ModuleList(
            [nn.ModuleList([Block() for _ in range(3)]) for _ in range(4)]
        )
        self.norm = nn.LayerNorm(768)
        st = torch.load(ckpt_path, map_location="cpu", weights_only=True)
        inner = st["teacher"]
        if isinstance(inner, dict) and "state_dict" in inner:
            inner = inner["state_dict"]
        own = self.state_dict()
        prefix = "backbone."
        missing = []
        for k in list(inner.keys()):
            if not k.startswith(prefix):
                inner.pop(k)
                continue
            name = k[len(prefix):]
            if name not in own:
                continue
            if tuple(inner[k].shape) == tuple(own[name].shape):
                own[name] = inner[k]
            else:
                missing.append((name, tuple(inner[k].shape), tuple(own[name].shape)))
        self.load_state_dict(own)
        if missing:
            print(f"mricore: {len(missing)} shape mismatches skipped "
                  f"(e.g. {missing[0]})", flush=True)
        if freeze:
            for p in self.parameters():
                p.requires_grad_(False)
        self.eval()

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        """[S, 3, H, W] -> [S, 197, 768] (CLS + 196 patch tokens)."""
        x = self.patch_embed(x).flatten(2).transpose(1, 2)  # [S, 196, 768]
        x = torch.cat([self.cls_token.expand(x.size(0), -1, -1), x], dim=1)
        x = x + self.pos_embed
        for group in self.blocks:
            for blk in group:
                x = blk(x)
        return self.norm(x)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.forward_features(x)