"""Differentiable loader for InternViT (OpenGVLab/InternViT-300M-448px-V2_5)."""
from __future__ import annotations

import glob
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModel

from _shim import apply_pretrained_shims

REPO = "OpenGVLab/InternViT-300M-448px-V2_5"
IMAGE_SIZE = 448
MEAN = (0.485, 0.456, 0.406)
STD = (0.229, 0.224, 0.225)
_STUB = Path(__file__).resolve().parent / "vendored" / "internvit" / "flash_attention.py"


def _patch_flash_attention_files() -> None:
    stub = _STUB.read_text()
    hub_path = hf_hub_download(REPO, "flash_attention.py")
    Path(hub_path).write_text(stub)
    for path in glob.glob(
        str(Path.home() / "workspace/hf-cache/modules/transformers_modules/**/flash_attention.py"),
        recursive=True,
    ):
        if "InternViT" in path or "InternViT_hyphen" in path:
            Path(path).write_text(stub)


def load_internvit(dtype: torch.dtype = torch.bfloat16) -> nn.Module:
    apply_pretrained_shims()
    _patch_flash_attention_files()
    config = AutoConfig.from_pretrained(REPO, trust_remote_code=True)
    config.use_flash_attn = False
    model = AutoModel.from_pretrained(
        REPO,
        config=config,
        trust_remote_code=True,
        dtype=dtype,
    )
    model.eval()
    model.requires_grad_(False)
    return model.to(device="cuda", dtype=dtype)


def image_feat_internvit(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
    """x: (1,3,H,W) in [0,1] on cuda, requires_grad=True. Returns (1, D) CLS feature."""
    assert x.shape[0] == 1 and x.dim() == 4
    mean = torch.tensor(MEAN, device=x.device, dtype=x.dtype).view(1, 3, 1, 1)
    std = torch.tensor(STD, device=x.device, dtype=x.dtype).view(1, 3, 1, 1)
    px = F.interpolate(x, size=(IMAGE_SIZE, IMAGE_SIZE), mode="bicubic", align_corners=False)
    px = (px - mean) / std
    dtype = next(model.parameters()).dtype
    out = model(pixel_values=px.to(dtype=dtype))
    feat = out.pooler_output.float()
    if feat.dim() == 1:
        feat = feat.unsqueeze(0)
    return feat