# %%
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

# %%
train_df = pd.read_csv("train.csv")
test_df = pd.read_csv("test.csv")
train_df.head()

# %%
class TrainDataset(Dataset):
    def __init__(self, path, transform):
        self.df = pd.read_csv(path)
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        img = Image.open(f"./images/{row['image_path']}").convert('RGB')
        if self.transform:
            img = self.transform(img)
        label = row['rotation_label']
        return img, label

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

# %%
train_ds = TrainDataset("./train.csv", transform=transform)
train = DataLoader(train_ds, batch_size=32, num_workers=0, shuffle=True)

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

# %%
weights = models.ResNet50_Weights.IMAGENET1K_V1
model = models.resnet50(weights=weights)
model.train()

model.fc = nn.Linear(2048,4)
model = model.to(device)
for param in model.parameters():
    param.requires_grad = False
for param in model.layer4.parameters():
    param.requires_grad = True
for param in model.fc.parameters():
    param.requires_grad = True

optimizer = optim.Adam([
    {'params': model.layer4.parameters(), 'lr': 1e-4},
    {'params': model.fc.parameters(), 'lr': 1e-3}
])
criterion = nn.CrossEntropyLoss()


# %%
for epoch in range(10):
    for img, label in train:
        img = img.to(device)
        label = label.to(device)
        optimizer.zero_grad()
        output = model(img)
        loss = criterion(output,label)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch+1}, loss {loss:.4f}")

# %%
class TestDataset(Dataset):
    def __init__(self, path, transform):
        self.df = pd.read_csv(path)
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        img = Image.open(f"./images/{row['image_path']}").convert('RGB')
        if self.transform:
            img = self.transform(img)
        
        return img

# %%
test_ds = TestDataset("./test.csv", transform=transform)
test = DataLoader(test_ds, batch_size=32, num_workers=0, shuffle=False)

# %%
test_df

# %%
pred=[]
model.eval()
for img in test:
    with torch.no_grad():
        img = img.to(device)
        pred_per_file = model(img)
        _, predicted = torch.max(pred_per_file, 1)
        pred.append(predicted.cpu())
        preds = torch.cat(pred)
preds

# %%
preds = preds.numpy()

# %%


# %%


# %%


# %%


# %%
map_ro = {0:0, 1:-90, 2:-180, 3:-270}

# %%
# corrected_images = []
# for idx, row in test_df.iterrows():
#     img = Image.open(f"./images/{row['image_path']}").convert('RGB')
#     angle = map_ro[preds[idx]]
#     img_corrected = img.rotate(angle)
#     corrected_images.append(img_corrected)

# %%
class CatalogDataset(Dataset):
    def __init__(self, path, transform):
        self.df = pd.read_csv(path)
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        img = Image.open(f"./images/{row['image_path']}").convert('RGB')
        if self.transform:
            img = self.transform(img)
        return img

# %%
catalog_ds = CatalogDataset("./vision_catalog.csv", transform=transform)
catalog = DataLoader(catalog_ds, batch_size=32, num_workers=0, shuffle=False)

# %%
encoder = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
encoder.fc = nn.Identity()
encoder.eval()
encoder = encoder.to(device)

# %%
catalog_embeddings = []
for img in catalog:
    with torch.no_grad():
        img = img.to(device)
        emb = encoder(img)
        catalog_embeddings.append(emb.cpu())
        
catalog_embeddings = torch.cat(catalog_embeddings)

# %%
class CorectedDataset(Dataset):
    def __init__(self, path, transform):
        self.df = pd.read_csv(path)
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        img = Image.open(f"./images/{row['image_path']}").convert('RGB')
        angle = map_ro[preds[index]]
        img_corrected = img.rotate(angle)
        if self.transform:
            img_corrected = self.transform(img_corrected)
        return img_corrected

# %%
Corected_ds = CorectedDataset("./test.csv", transform=transform)
corrected = DataLoader(Corected_ds, batch_size=32, num_workers=0, shuffle=False)

# %%
test_embeddings = []
for img in corrected:
    with torch.no_grad():
        img = img.to(device)
        emb = encoder(img)
        test_embeddings.append(emb.cpu())
        
test_embeddings = torch.cat(test_embeddings)

# %%
catalog_norm = catalog_embeddings / catalog_embeddings.norm(dim=1, keepdim=True)
test_norm = test_embeddings / test_embeddings.norm(dim=1, keepdim=True)
scores = test_norm @ catalog_norm.T
best_indices = scores.argmax(dim=1)

# %%
catalog_df = pd.read_csv('vision_catalog.csv')
predicted_gallery_ids=[]
for i in range(len(best_indices)):
    predicted_gallery_ids.append(catalog_df['gallery_id'].values[best_indices[i]])
predicted_gallery_ids

# %%
rows = []
for idx, row in test_df.iterrows():
    rows.append({"subtaskID":1, "datapointID":row['datapoint_id'],'answer':preds[idx]})
    rows.append({"subtaskID":2, "datapointID":row['datapoint_id'],'answer':predicted_gallery_ids[idx]})
pd.DataFrame(rows).to_csv('subs.csv',index=False)

# %%



