from functools import lru_cache
from typing import Any
import open_clip
from PIL import Image
import numpy as np


class VisionLanguageModel:
    """
    A class that wraps the OpenCLIP model for vision-language tasks.

    Args:
        model_name (str, optional): The name of the OpenCLIP model. Defaults to "ViT-L-14".
        pretrained (str, optional): The name of the pretrained model. Defaults to "laion2B-s32B-b82K".
        device (str, optional): The device to use for computations. Defaults to "cuda".

    Attributes:
        device (str): The device used for computations.
        model: The OpenCLIP model instance.
        processor: The image processor instance.
        tokenizer: The tokenizer instance.
    """

    def __init__(
        self,
        model_name: str = "ViT-B-16",
        pretrained: str = "laion2B-s34B-b88K",
        device: str = "cuda",
    ):
        """
        Initializes the VisionLanguageModel instance.

        Args:
            model_name (str, optional): The name of the OpenCLIP model. Defaults to "ViT-B-16".
            pretrained (str, optional): The name of the pretrained model. Defaults to "laion2B-s34B-b88K".
            device (str, optional): The device to use for computations. Defaults to "cuda".
        """
        self.device: str = device
        self._load_model(model_name, pretrained)

    @lru_cache(maxsize=1)
    def _load_model(self, model_name: str, pretrained: str):
        """
        Loads the OpenCLIP model and tokenizer.

        Args:
            model_name (str): The name of the OpenCLIP model.
            pretrained (str): The name of the pretrained model.
        """
        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
            model_name=model_name, pretrained=pretrained, device=self.device
        )
        self.model.eval()
        self.tokenizer = open_clip.get_tokenizer(model_name=model_name)

    def get_embedding(self, input: str | Image.Image | Any) -> np.ndarray:
        """
        Computes the embedding for the given input.

        Args:
            input (str | Image.Image | Any): The input to compute the embedding for.
                Can be a string, a PIL Image, or any other type (which will raise a ValueError).

        Returns:
            np.ndarray: The computed embedding as a NumPy array.

        Raises:
            ValueError: If the input type is not supported.

        Examples:
            >>> model = VisionLanguageModel()
            >>> text_embedding = model.get_embedding("This is a sample text")
            >>> image = Image.open("image.jpg").convert("RGB")
            >>> image_embedding = model.get_embedding(image)
        """
        if isinstance(input, str):
            tokens = self.tokenizer(input).to(self.device)
            vector = self.model.encode_text(tokens)
        elif isinstance(input, Image.Image):
            image_input = self.preprocess(input).unsqueeze(0).to(self.device)
            vector = self.model.encode_image(image_input)
        else:
            raise ValueError(f"Invalid input type: {type(input)}")

        vector /= vector.norm(dim=-1, keepdim=True)
        return vector.cpu().detach().numpy().astype("float32")