# -- Code Cell --
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 pandas as pd
import os
from sklearn.decomposition import PCA

# -- Code Cell --
class customDataset(Dataset):
    def __init__(self, path, transform):
        self.transform=transform
        self.path = path
        self.files = os.listdir(path)
    def __len__(self):
        return len(self.files)
    def __getitem__(self, index):
        img = Image.open(f"{self.path}/img_{index}.png").convert('RGB')
        if self.transform:
            img = self.transform(img)
        return img

# -- Code Cell --
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# -- Code Cell --
train_ds = customDataset("./train", transform=transform)
train = DataLoader(train_ds, batch_size=32, shuffle=False, num_workers=0)

# -- Code Cell --
device = torch.device('cuda')

# -- Code Cell --
vgg = torch.load("vgg19.pth", weights_only=False).to(device)
vgg.eval()
vgg_extractor = nn.Sequential(*list(vgg.children())[:-1]).to(device)

# -- Code Cell --
def extract_features(loader):
    all_features = []
    with torch.no_grad():
        for imgs in loader:
            imgs = imgs.to(device)
            features = vgg_extractor(imgs)
            features = features.view(features.size(0), -1)   
            all_features.append(features.cpu())
    return torch.cat(all_features, dim=0)

# -- Code Cell --
test_ds = customDataset("./test", transform=transform)
test = DataLoader(test_ds, batch_size=32, shuffle=False, num_workers=0)

# -- Code Cell --
train_X = extract_features(train)
test_X = extract_features(test)

# -- Code Cell --
import umap
from sklearn.cluster import KMeans

# -- Code Cell --
X = test_X.numpy()
emb = umap.UMAP(n_neighbors=10,n_components=30,random_state=42).fit_transform(X)
kmeans = KMeans(n_clusters=2, random_state=42)
preds = kmeans.fit_predict(emb)

preds = torch.tensor(preds)

# -- Code Cell --
preds = 1-preds

# -- Code Cell --
# pca = PCA(n_components=32)
# pca.fit(train_X.numpy())

# train_reduced = pca.transform(train_X.numpy())
# test_reduced = pca.transform(test_X.numpy())

# train_reduced = torch.tensor(train_reduced, dtype=torch.float32)
# test_reduced = torch.tensor(test_reduced, dtype=torch.float32)

# centroid = train_reduced.mean(dim=0)
# train_distances = torch.norm(train_reduced - centroid, dim=1)
# test_distances = torch.norm(test_reduced - centroid, dim=1)

# -- Code Cell --
# centroid = train_reduced.mean(dim=0)
# train_distances = torch.norm(train_reduced - centroid, dim=1)
# test_distances = torch.norm(test_reduced - centroid, dim=1)

# -- Code Cell --
# distances = torch.norm(test_reduced - centroid, dim=1)
# preds = (distances > 12).long()
# print(f"reale: {(preds==0).sum()}, fake: {(preds==1).sum()}")

# -- Code Cell --
rows = []
for i in range(len(preds)):
    rows.append({'subtaskID': 1, 'datapointID': i, 'answer': int(preds[i])})

pd.DataFrame(rows).to_csv('submission.csv', index=False)

# -- Code Cell --
