import torch
from PIL import Image
from torchvision import transforms
from transformers import CLIPModel as HFCLIPModel, CLIPConfig


class BioCLIPv2Model(HFCLIPModel):
    def __init__(self):
        super().__init__(self._build_config())

    @staticmethod
    def _build_config():
        return CLIPConfig(
            projection_dim=768,
            logit_scale_init_value=2.6592,
            text_config={
                "hidden_size": 768,
                "intermediate_size": 3072,
                "num_attention_heads": 12,
                "num_hidden_layers": 12,
                "max_position_embeddings": 77,
                "vocab_size": 49408,
                "hidden_act": "gelu",
                "layer_norm_eps": 1e-5,
                "projection_dim": 768,
                "bos_token_id": 0,
                "eos_token_id": 2,
                "pad_token_id": 1,
            },
            vision_config={
                "hidden_size": 1024,
                "intermediate_size": 4096,
                "num_attention_heads": 16,
                "num_hidden_layers": 24,
                "image_size": 224,
                "patch_size": 14,
                "hidden_act": "gelu",
                "layer_norm_eps": 1e-5,
                "projection_dim": 768,
            },
        )


class BioCLIPv2Compute(torch.nn.Module):
    MEAN = (0.48145466, 0.4578275, 0.40821073)
    STD = (0.26862954, 0.26130258, 0.27577711)

    def __init__(self, model: BioCLIPv2Model, tokenizer, max_length: int = 77):
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.image_transform = transforms.Compose([
            transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(self.MEAN, self.STD),
        ])

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

    @property
    def dtype(self):
        return next(self.model.parameters()).dtype

    def _preprocess_images(self, images):
        if isinstance(images, Image.Image):
            images = [images]
        images = [img.convert("RGB") for img in images]
        pixel_values = torch.stack([self.image_transform(img) for img in images])
        return pixel_values.to(device=self.device, dtype=self.dtype)

    def _tokenize(self, text):
        if isinstance(text, str):
            text = [text]
        tokens = self.tokenizer(
            text, padding=True, truncation=True,
            max_length=self.max_length, return_tensors="pt",
        )
        return {k: v.to(self.device) for k, v in tokens.items()}

    @torch.no_grad()
    def get_image_features(self, images):
        pixel_values = self._preprocess_images(images)
        features = self.model.get_image_features(pixel_values=pixel_values)
        return torch.nn.functional.normalize(features, dim=-1)

    @torch.no_grad()
    def get_text_features(self, text):
        tokens = self._tokenize(text)
        features = self.model.get_text_features(**tokens)
        return torch.nn.functional.normalize(features, dim=-1)

    @torch.no_grad()
    def forward(self, text: str | list[str], images):
        if isinstance(text, str):
            text = [text]
        if isinstance(images, Image.Image):
            images = [images]
        if len(text) == 1 and len(images) > 1:
            text = text * len(images)
        if len(images) == 1 and len(text) > 1:
            images = images * len(text)
        image_features = self.get_image_features(images)
        text_features = self.get_text_features(text)
        scores = (text_features * image_features).sum(dim=-1)
        scores = self.model.logit_scale.exp() * scores
        return scores

    @torch.no_grad()
    def similarity_matrix(self, text: str | list[str], images):
        image_features = self.get_image_features(images)
        text_features = self.get_text_features(text)
        scores = text_features @ image_features.T
        scores = self.model.logit_scale.exp() * scores
        return scores