from typing import Union
import numpy as np
from PIL import Image
import torch
from transformers import CLIPProcessor, CLIPModel
from .base import BaseEncoder


class ImageEncoder(BaseEncoder):
    """Image encoder using CLIP."""

    def __init__(self, model_name: str = "openai/clip-vit-base-patch32", revision: str = "main"):
        self.model_name = model_name
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        # The revision parameter remains caller-configurable to support
        # pinning trusted model versions.  Bandit is explicitly informed that
        # the revision is supplied.
        self.model = CLIPModel.from_pretrained(
            model_name, revision=revision).to(self.device)  # nosec B615
        self.processor = CLIPProcessor.from_pretrained(
            model_name, revision=revision)  # nosec B615

    def encode(self, data: Union[str, Image.Image]) -> np.ndarray:
        """
        Encode image into a vector.

        Args:
            data: File path to image or PIL Image object

        Returns:
            Numpy array representing the embedding vector
        """
        if isinstance(data, str):
            image = Image.open(data)
        else:
            image = data

        inputs = self.processor(
            images=image,
            return_tensors="pt").to(
            self.device)

        with torch.no_grad():
            image_features = self.model.get_image_features(**inputs)

        # Normalize features
        image_features = image_features / \
            image_features.norm(p=2, dim=-1, keepdim=True)

        return image_features.cpu().numpy().flatten()

    @property
    def dimension(self) -> int:
        """Return the dimension of the embedding vector."""
        return self.model.config.projection_dim
