# -- 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 --
model.config.return_loss = True

# -- Code Cell --
from torch.utils.data import Dataset
from PIL import Image

class CLIPDataset(Dataset):
    def __init__(self, df, processor):
        self.df = df.reset_index(drop=True)
        self.processor = processor

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        image = Image.open(row["image_path"]).convert("RGB")

        correct = row["correct_words"]
        positive_text = correct[0] if correct else row["word_choices"][0]

        inputs = self.processor(
            images=image,
            text=positive_text,
            return_tensors="pt",
            padding="max_length",
            truncation=True,
            max_length=77,
        )

        return {
            "input_ids": inputs["input_ids"].squeeze(0),
            "attention_mask": inputs["attention_mask"].squeeze(0),
            "pixel_values": inputs["pixel_values"].squeeze(0),
            "labels": torch.tensor(0)  
        }
dataset = CLIPDataset(df_train, processor)

# -- Code Cell --
class CLIPWithLoss(CLIPModel):
    def forward(self, input_ids=None, pixel_values=None, attention_mask=None, **kwargs):
        outputs = super().forward(
            input_ids=input_ids,
            pixel_values=pixel_values,
            attention_mask=attention_mask,
            return_loss=False,  
            **kwargs,
        )

        img_feats = F.normalize(outputs.image_embeds, dim=-1)
        txt_feats = F.normalize(outputs.text_embeds, dim=-1)

        logits = img_feats @ txt_feats.T 
        labels = torch.arange(len(logits), device=logits.device)
        loss = (F.cross_entropy(logits, labels) + F.cross_entropy(logits.T, labels)) / 2

        outputs.loss = loss
        return outputs

# -- Code Cell --
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
from tqdm.notebook import tqdm

def make_collate(processor):
    def collate_fn(batch):
        images, pos_texts, all_choices = zip(*batch)
        img_inputs = processor(images=list(images), return_tensors="pt", padding=True)
        txt_inputs = processor(text=list(pos_texts), return_tensors="pt",
                               padding=True, truncation=True)
        return img_inputs, txt_inputs, all_choices
    return collate_fn

def contrastive_loss(image_features, text_features, temperature=0.07):
    """Standard symmetric CLIP contrastive loss."""
    logits = (image_features @ text_features.T) / temperature
    labels = torch.arange(len(logits), device=logits.device)
    loss_i = F.cross_entropy(logits, labels)
    loss_t = F.cross_entropy(logits.T, labels)
    return (loss_i + loss_t) / 2

def finetune_clip(df_train, model, processor, epochs=5, lr=1e-5, batch_size=8):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    for name, param in model.named_parameters():
        param.requires_grad = False
    for name, param in model.named_parameters():
        if any(k in name for k in [
            "visual_projection", "text_projection",
            "vision_model.encoder.layers.23",
            "text_model.encoder.layers.11",
        ]):
            param.requires_grad = True

    dataset = CLIPDataset(df_train, processor)
    loader  = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad], lr=lr
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=epochs * len(loader)
    )

    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for batch in tqdm(loader, desc=f"Epoch {epoch+1}"): 
            img_inputs = {"pixel_values": batch["pixel_values"].to(device)}
            txt_inputs = {
                "input_ids": batch["input_ids"].to(device),
                "attention_mask": batch["attention_mask"].to(device),
            }

            img_feats = F.normalize(model.get_image_features(**img_inputs), dim=-1)
            txt_feats = F.normalize(model.get_text_features(**txt_inputs), dim=-1)

            loss = contrastive_loss(img_feats, txt_feats)
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            total_loss += loss.item()

        print(f"Epoch {epoch+1} — avg loss: {total_loss/len(loader):.4f}")

    return model

# ── Run it ────────────────────────────────────────────────────────────────────
model = finetune_clip(df_train, model, processor, epochs=10, lr=1e-5, batch_size=8)

# -- Code Cell --
model = model.to("cpu")

# -- 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_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("subiv2.csv", index= False)

# -- Code Cell --
