# %%
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms
from PIL import Image
import os

# %%
class TrainDataset(Dataset):
    def __init__(self, transform):
        self.transform = transform
    def __len__(self):
        return len(os.listdir("./train/"))
    def __getitem__(self, index):
        img = Image.open(f"./train/img_{index}.png")
        if self.transform:
            img = self.transform(img)
        return img

# %%
transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# %%
train_ds = TrainDataset(transform)
train_loader = DataLoader(train_ds,num_workers=0,batch_size=64,shuffle=False)

# %%
device = torch.device('cuda')

# %%
model = torch.load("resnet152.pth",weights_only=False)
model.eval()
model.fc = nn.Identity()
model = model.to(device)

# %%
embeddings = []
with torch.no_grad():
    for img in train_loader:
        img = img.to(device)
        emb = model(img)
        embeddings.append(emb.cpu())
embeddings = torch.cat(embeddings)
embeddings = embeddings.numpy()

# %%
embeddings

# %%
class TestDataset(Dataset):
    def __init__(self, transform):
        self.transform = transform
    def __len__(self):
        return len(os.listdir("./test/"))
    def __getitem__(self, index):
        img = Image.open(f"./test/img_{index}.png")
        if self.transform:
            img = self.transform(img)
        return img

# %%
test_ds = TestDataset(transform)
test_loader = DataLoader(test_ds,num_workers=0,batch_size=64,shuffle=False)

# %%
embeddings_test = []
with torch.no_grad():
    for img in test_loader:
        img = img.to(device)
        emb = model(img)
        embeddings_test.append(emb.cpu())
embeddings_test = torch.cat(embeddings_test)
embeddings_test = embeddings_test.numpy()

# %%
embeddings_test

# %%
#STRATEGIE: COSINE SIMILARITY >0.9 == NU E GHOST

# %%
import numpy as np

# %%
pred_final = []
for test_img in range(len(embeddings_test)):
    values_per_test_img = []
    for train_img in range(len(embeddings)):
        embeddings_norm = embeddings[train_img] / np.linalg.norm(embeddings[train_img], axis=0, keepdims=True)
        embeddings_test_norm = embeddings_test[test_img] / np.linalg.norm(embeddings_test[test_img], axis=0, keepdims=True)
        cos_sim = embeddings_norm @ embeddings_test_norm.T
        # print(nn.functional.cosine_similarity(torch.tensor(embeddings[train_img]), torch.tensor(embeddings_test[test_img])))
        # print(f"imaginea din train: img_{train_img}.png cu imaginea din test: img_{test_img}.png au similaritate: {cos_sim}")
        values_per_test_img.append(cos_sim)
    score = np.max(values_per_test_img)
    # print(f"test img_{test_img}.png are mean cos: {score}")
    pred_final.append(1 if score>0.88 else 0)

# %%
pred_final

# %%
import pandas as pd

# %%
id = []
for i in range(2000):
    id.append(i)
id

# %%
subs1 = pd.DataFrame({
    "datapointID":id,
    "subtaskID":1,
    'answer':pred_final
}).to_csv('subs.csv',index=False)

# %%



