# -- Code Cell --
import pandas as pd
import os

# -- Code Cell --
X_train = []
y_json = []
X_test = []
y_tst_json = []
id = []
for folders in os.listdir("./train"):
    folder_path = os.path.join("./train", folders)
    for pth in os.listdir(folder_path):
        if pth.endswith(".png"):
            X_train.append(os.path.join(folder_path, pth))
        if pth.endswith(".json"):
            y_json.append(os.path.join(folder_path, pth))
for folders in os.listdir("./test"):
    folder_path = os.path.join("./test", folders)
    id.append(folders)
    for pth in os.listdir(folder_path):
        if pth.endswith(".png"):
            X_test.append(os.path.join(folder_path, pth))
        if pth.endswith(".json"):
            y_tst_json.append(os.path.join(folder_path, pth))

# -- Code Cell --
import json
import pandas as pd
df_train = pd.DataFrame({})
for i in range(len(y_json)):
    with open(y_json[i]) as f:
        data = json.load(f)
    df = pd.json_normalize(data)
    df_train = pd.concat([df_train, df])
df_train["image_path"] = X_train
df_train = df_train[["image_path", "word_choices", "correct_words"]]

# -- Code Cell --
import json
import pandas as pd
df_test = pd.DataFrame({})
for i in range(len(y_tst_json)):
    with open(y_tst_json[i]) as f:
        data = json.load(f)
    df = pd.json_normalize(data)
    df_test = pd.concat([df_test, df])
df_test["image_path"] = X_test


# -- Code Cell --
df_test = df_test[["image_path", "word_choices"]]

# -- Code Cell --
df_train["word_choices"] = df_train["word_choices"].apply(lambda x: [f"a photo of a {_}" for _ in x])

# -- Code Cell --
df_test["word_choices"] = df_test["word_choices"].apply(lambda x: [f"a photo of a {_}" for _ in x])

# -- Code Cell --
df_train["correct_words"] = df_train["correct_words"].apply(lambda x: [f"a photo of a {_}" for _ in x])

# -- Code Cell --
df_test["word_choices"] = df_test["word_choices"].apply(lambda x: [f"a photo of a {_}" for _ in x])

# -- Code Cell --
from transformers import CLIPProcessor, CLIPModel

model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")

# -- Code Cell --
import torch
import torch.nn.functional as F
from PIL import Image
from tqdm.notebook import tqdm
text_embeddings = []
image_embeddings = []
for i in tqdm(range(0, len(df_test))):
    #TEXT EMBEDDINGS
    batch = df_test["word_choices"].iloc[i]

    inputs = processor(
        text=batch,
        return_tensors="pt",
        padding=True,
        truncation=True
    )

    with torch.no_grad():
        text_features = model.get_text_features(**inputs)

    text_features = F.normalize(text_features, dim=-1)
    text_embeddings.append(text_features)

    #IMAGE EMBEDDINGS
    img = Image.open(df_test["image_path"].iloc[i]).convert("RGB")
    inputs = processor(images= img, return_tensors = "pt")
    with torch.no_grad():
        image_features = model.get_image_features(**inputs)
    image_features = F.normalize(image_features, dim = -1).squeeze(0)
    image_embeddings.append(image_features)
    print
# concatenate everything
text_embeddings = torch.stack(text_embeddings, dim=0)
image_embeddings = torch.stack(image_embeddings, dim=0)

# -- Code Cell --
image_embeddings.shape

# -- Code Cell --
image_expanded = image_embeddings.unsqueeze(1)  

similarity = (image_expanded * text_embeddings).sum(dim=-1)


# -- Code Cell --
topk = similarity.topk(k=5, dim=1)
top_indices = topk.indices  # shape: (200, 3)

# -- Code Cell --
predictions = [
    [
        df_test["word_choices"].iloc[i][idx.item()].split(' ')[-1]
        for idx in top_indices[i]
    ]
    for i in range(top_indices.shape[0])
]

# -- Code Cell --
pd.DataFrame({
    "ID" : id,
    "word1" :[x[0] for x in predictions],
    "word2" : [x[1] for x in predictions],
    "word3" : [x[2] for x in predictions],
    "word4" : [x[3] for x in predictions],
    "word5" : [x[4] for x in predictions],
}).to_csv("subi.csv", index= False)