from transformers import DistilBertTokenizer
import torch
from utils import *
from clip import CLIPModel
from tqdm import tqdm
from config import CFG
import matplotlib.pyplot as plt
import cv2
import torch.nn.functional as F
from data import *

img_embeddings = None
def get_image_embeddings(valid_df, model_path, get_emb = False):
    tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
    valid_loader = get_dataset_loader(valid_df, tokenizer, mode="valid")
    
    model = CLIPModel().to(CFG.device)
    model.load_state_dict(torch.load(model_path, map_location=CFG.device))
    model.eval()
    
    valid_image_embeddings = []
    if get_emb:
        with torch.no_grad():
            for batch in tqdm(valid_loader):
                image_features = model.image_encoder(batch["image"].to(CFG.device))
                image_embeddings = model.image_projection(image_features)
                valid_image_embeddings.append(image_embeddings)
        torch.save(torch.cat(valid_image_embeddings), CFG.img_emb_path)
    return model

def load_img_embeddings():
    global img_embeddings
    img_embeddings = torch.load(CFG.img_emb_path, map_location=CFG.device)
    img_embeddings = F.normalize(img_embeddings, p=2, dim=-1)

def find_matches(model, query, image_filenames, n=9):
    tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
    encoded_query = tokenizer([query])
    batch = {
        key: torch.tensor(values).to(CFG.device)
        for key, values in encoded_query.items()
    }
    with torch.no_grad():
        text_features = model.text_encoder(
            input_ids=batch["input_ids"], attention_mask=batch["attention_mask"]
        )
        text_embeddings = model.text_projection(text_features)
    
    image_embeddings_n = img_embeddings
    text_embeddings_n = F.normalize(text_embeddings, p=2, dim=-1)
    dot_similarity = text_embeddings_n @ image_embeddings_n.T
    
    values, indices = torch.topk(dot_similarity.squeeze(0), n * 5)
    matches = [image_filenames[idx] for idx in indices[::5]]
    plt.clf()
    plt.cla()
    fig, axes = plt.subplots(3, 3, figsize=(10, 10))
    for match, ax in zip(matches, axes.flatten()):
        image = cv2.imread(f"{CFG.image_path}/{match}")
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        ax.imshow(image)
        ax.axis("off")
    return fig

load_img_embeddings()
##_, valid_df = gen_train_valid_dfs()
#model = get_image_embeddings(valid_df, "./model/best.pt")
#find_matches(model, query="boy is playing", image_filenames=valid_df['image'].values, n=9)

