import functools
import logging
import typing

import beartype
import torch
from jaxtyping import Float, jaxtyped
from torch import Tensor
from torchvision.transforms import v2

logger = logging.getLogger("modeling.py")


@jaxtyped(typechecker=beartype.beartype)
class SplitDinov2(torch.nn.Module):
    def __init__(self, *, split_at: int):
        super().__init__()

        self.vit = torch.hub.load("facebookresearch/dinov2", "dinov2_vitb14_reg").eval()
        self.split_at = split_at

    def forward_start(
        self, x: Float[Tensor, "batch channels width height"]
    ) -> Float[Tensor, "batch total_patches dim"]:
        x_BPD = self.vit.prepare_tokens_with_masks(x)
        for blk in self.vit.blocks[: self.split_at]:
            x_BPD = blk(x_BPD)

        return x_BPD

    def forward_end(
        self, x_BPD: Float[Tensor, "batch total_patches dim"]
    ) -> Float[Tensor, "batch patches dim"]:
        for blk in self.vit.blocks[-self.split_at :]:
            x_BPD = blk(x_BPD)

        x_BPD = self.vit.norm(x_BPD)
        return x_BPD[:, self.vit.num_register_tokens + 1 :]


@functools.cache
def load_vit(device: str) -> tuple[SplitDinov2, typing.Callable]:
    vit = SplitDinov2(split_at=11).to(device)
    vit_transform = v2.Compose([
        v2.Resize(size=(256, 256)),
        v2.CenterCrop(size=(224, 224)),
        v2.ToImage(),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize(mean=[0.4850, 0.4560, 0.4060], std=[0.2290, 0.2240, 0.2250]),
    ])
    logger.info("Loaded ViT.")

    return vit, vit_transform