import torch
from einops import rearrange
from jaxtyping import Float
from PIL import Image
from torch import Tensor
from torch import nn
from transformers import AutoImageProcessor
from transformers import AutoModel
from transformers.feature_extraction_utils import BatchFeature


__version__ = "0.1.0"

TypeClsToken = Float[Tensor, "batch_size embed_dim"]
TypePatchTokensFlat = Float[Tensor, "batch_size (height width) embed_dim"]
TypePatchTokens = Float[Tensor, "batch_size embed_dim height width"]
TypeInputImages = Image.Image | list[Image.Image]


class RadDino(nn.Module):
    _REPO = "microsoft/rad-dino"

    def __init__(self):
        super().__init__()
        self.model = AutoModel.from_pretrained(self._REPO).eval()
        self.processor = AutoImageProcessor.from_pretrained(self._REPO, use_fast=False)

    @property
    def device(self) -> torch.device:
        return next(self.model.parameters()).device

    def preprocess(self, image_or_images: TypeInputImages) -> BatchFeature:
        return self.processor(image_or_images, return_tensors="pt")

    def encode(self, inputs: BatchFeature) -> tuple[TypeClsToken, TypePatchTokensFlat]:
        outputs = self.model(**inputs)
        cls_token = outputs.last_hidden_state[:, 0]
        patch_tokens = outputs.last_hidden_state[:, 1:]
        return cls_token, patch_tokens

    def reshape_patch_tokens(
        self,
        patch_tokens_flat: TypePatchTokensFlat,
    ) -> TypePatchTokens:
        input_size = self.processor.crop_size["height"]
        patch_size = self.model.config.patch_size
        embeddings_size = input_size // patch_size
        patches_grid = rearrange(
            patch_tokens_flat,
            "batch (height width) embed_dim -> batch embed_dim height width",
            height=embeddings_size,
        )
        return patches_grid

    @torch.inference_mode()
    def extract_features(
        self,
        image_or_images: TypeInputImages,
    ) -> tuple[TypeClsToken, TypePatchTokens]:
        inputs = self.preprocess(image_or_images).to(self.device)
        cls_token, patch_tokens_flat = self.encode(inputs)
        patch_tokens = self.reshape_patch_tokens(patch_tokens_flat)
        return cls_token, patch_tokens

    def extract_cls_token(self, image_or_images: TypeInputImages) -> TypeClsToken:
        cls_token, _ = self.extract_features(image_or_images)
        return cls_token

    def extract_patch_tokens(self, image_or_images: TypeInputImages) -> TypePatchTokens:
        _, patch_tokens = self.extract_features(image_or_images)
        return patch_tokens

    def forward(self, *args) -> tuple[TypeClsToken, TypePatchTokens]:
        return self.extract_features(*args)