import clip  
import torch

from numpy import ndarray  
from typing import List  
from PIL import Image 

class ClipEmbeddings:

    def __init__(self, model_name: str = "ViT-B/32", device: str = "cpu"):
        
        self.device = device  # Store the specified device for model execution
        
        self.model, self.preprocess = clip.load(model_name, self.device)

    def __call__(self, docs: List[str]) -> List[ndarray]:
        # Define a method that takes a list of image file paths (docs) as input
        list_of_embeddings = []  # Create an empty list to store the image embeddings
        for image_path in docs:
            image = Image.open(image_path)  # Open and load an image from the provided path
            
            image = image.resize((224, 224))  
            # Preprocess the image and move it to the specified device
            image_input = self.preprocess(image).unsqueeze(0).to(self.device)  
            with torch.no_grad():
                # Compute the image embeddings using the CLIP model and convert 
                #them to NumPy arrays
                embeddings = self.model.encode_image(image_input).cpu().detach().numpy()
            list_of_embeddings.append(list(embeddings[0])) 
        return list_of_embeddings  

    def get_text_embeddings(self, text: str) -> List[ndarray]:
        # Define a method that takes a text string as input
        text_token = clip.tokenize(text)  # Tokenize the input text
        with torch.no_grad():
            # Compute the text embeddings using the CLIP model and convert them to NumPy arrays
            text_embeddings = self.model.encode_text(text_token).cpu().detach().numpy()
        return list(text_embeddings[0])