"""Differentiable loader for MoonViT (moonshotai/MoonViT-SO-400M)."""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel

from _shim import apply_pretrained_shims

REPO = "moonshotai/MoonViT-SO-400M"
PATCH_SIZE = 14
DEFAULT_SIZE = 448


def load_moonvit(dtype: torch.dtype = torch.bfloat16) -> nn.Module:
    apply_pretrained_shims()
    model = AutoModel.from_pretrained(
        REPO,
        trust_remote_code=True,
        attn_implementation="sdpa",
        dtype=dtype,
    )
    model.eval()
    model.requires_grad_(False)
    return model.to(device="cuda", dtype=dtype)


def _patches_from_image(x: torch.Tensor, patch_size: int = PATCH_SIZE) -> tuple[torch.Tensor, torch.Tensor]:
    """Unfold (1,3,H,W) float image into native-resolution patch batch (N,3,ps,ps)."""
    _, _, h, w = x.shape
    assert h % patch_size == 0 and w % patch_size == 0, f"H,W must be divisible by {patch_size}"
    gh, gw = h // patch_size, w // patch_size
    patches = x.unfold(2, patch_size, patch_size).unfold(3, patch_size, patch_size)
    patches = patches.permute(0, 2, 4, 1, 3, 5).reshape(gh * gw, 3, patch_size, patch_size)
    grid = torch.tensor([[gh, gw]], device=x.device, dtype=torch.int32)
    return patches, grid


def image_feat_moonvit(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
    """x: (1,3,H,W) in [0,1] on cuda, requires_grad=True. Returns (1, D) pooled feature."""
    assert x.shape[0] == 1 and x.dim() == 4
    if x.shape[-1] != DEFAULT_SIZE or x.shape[-2] != DEFAULT_SIZE:
        x = F.interpolate(x, size=(DEFAULT_SIZE, DEFAULT_SIZE), mode="bicubic", align_corners=False)
    dtype = next(model.parameters()).dtype
    patches, grid = _patches_from_image(x)
    tokens_list = model(patches.to(dtype=dtype), grid)
    tokens = tokens_list[0].float()
    # tokens: (num_merged_positions, merge_kernel^2, hidden)
    feat = tokens.reshape(-1, tokens.shape[-1]).mean(dim=0, keepdim=True)
    return feat