import clip

import tqdm

import numpy as np

import torch

from reproducibility.embedders.internal_datasets import *

from torch.utils.data import DataLoader

from reproducibility.utils.cacher import cache_hit_or_miss, cache_numpy_object, cache_hit_or_miss_raw_filename, cache_numpy_object_raw_filename

class CLIPEmbedder:

    def __init__(self, model, preprocess, name, backbone):

        self.model = model

        self.preprocess = preprocess

        self.name = name

        self.backbone = backbone

    def image_embedder(self, list_of_images, device="cuda", num_workers=1, batch_size=32, additional_cache_name=""):

        hit_or_miss = cache_hit_or_miss_raw_filename(self.name + "img" + additional_cache_name, self.backbone)

        if hit_or_miss is not None:

            return hit_or_miss

        else:

            hit = self.embed_images(list_of_images, device=device, num_workers=num_workers, batch_size=batch_size)

            cache_numpy_object_raw_filename(hit, self.name + "img" + additional_cache_name, self.backbone)

            return hit

    def text_embedder(self, list_of_labels, device="cuda", num_workers=1, batch_size=32, additional_cache_name=""):

        hit_or_miss = cache_hit_or_miss(self.name + "txt" + additional_cache_name, self.backbone)

        if hit_or_miss is not None:

            return hit_or_miss

        else:

            hit = self.embed_text(list_of_labels, device=device, num_workers=num_workers, batch_size=batch_size)

            cache_numpy_object(hit, self.name + "txt" + additional_cache_name, self.backbone)

            return hit

    def embed_images(self, list_of_images, device="cuda", num_workers=1, batch_size=32):

        train_dataset = CLIPImageDataset(list_of_images, self.preprocess)

        dataloader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers)

        image_embeddings = []

        total = len(list_of_images) // batch_size

        pbar = tqdm.tqdm(total=total, position=0)

        with torch.no_grad():

            for images in dataloader:

                images = images.to(device)

                image_embeddings.extend(self.model.encode_image(images).detach().cpu().numpy())

                pbar.update(1)

            pbar.close()

        image_embeddings = np.array(image_embeddings)

        image_embeddings = image_embeddings / np.linalg.norm(image_embeddings, axis=1, keepdims=True)

        return image_embeddings

    def embed_text(self, list_of_labels, device="cuda", num_workers=1, batch_size=32):

        train_dataset = CLIPCaptioningDataset(list_of_labels)

        dataloader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers)

        text_embeddings = []

        total = len(list_of_labels) // batch_size

        pbar = tqdm.tqdm(total=total, position=0)

        with torch.no_grad():

            for captions in dataloader:

                idx = clip.tokenize(captions, truncate=True).to(device)

                text_embeddings.extend(self.model.encode_text(idx).detach().cpu().numpy())

                pbar.update(1)

            pbar.close()

        text_embeddings = np.array(text_embeddings)

        text_embeddings = text_embeddings / np.linalg.norm(text_embeddings, axis=1, keepdims=True)

        return text_embeddings