from lexy.models.document import DocumentBase
from lexy.transformers import lexy_transformer

import torch
from PIL.Image import Image
from transformers import CLIPProcessor, CLIPModel

torch.set_num_threads(1)
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# move model to device if possible
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)


@lexy_transformer(name="image.embeddings.clip")
def image_embeddings_clip(
    images: Image | DocumentBase | list[Image | DocumentBase],
) -> torch.Tensor:
    """Embed images using CLIP.

    Args:
        images: A single image or DocumentBase instance, or a list of images or
            DocumentBase instances to embed.

    Returns:
        torch.Tensor: The embeddings of the provided images.
    """

    if isinstance(images, DocumentBase):
        images = images.image
    elif isinstance(images, list):
        images = [i.image if isinstance(i, DocumentBase) else i for i in images]

    image_batch = processor(text=None, images=images, return_tensors="pt")[
        "pixel_values"
    ].to(device)

    if isinstance(images, list):
        return model.get_image_features(image_batch)
    else:
        return model.get_image_features(image_batch)[0]


@lexy_transformer(name="text.embeddings.clip")
def text_embeddings_clip(
    text: list[str | DocumentBase] | str | DocumentBase,
) -> torch.Tensor:
    """Embed text using CLIP.

    Args:
        text: A single string or DocumentBase instance, or a list of strings or
            DocumentBase instances to embed.

    Returns:
        torch.Tensor: The embeddings of the provided text.
    """
    if isinstance(text, DocumentBase):
        text = text.content
    elif isinstance(text, list):
        text = [s.content if isinstance(s, DocumentBase) else s for s in text]
    tokens = processor(text=text, padding=True, images=None, return_tensors="pt").to(
        device
    )
    text_emb = model.get_text_features(**tokens)
    if isinstance(text, list):
        return text_emb
    else:
        return text_emb[0]
