# -- 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
from torchvision.transforms import functional as TF
import cv2 as cv
import numpy as np

# -- Code Cell --
import pandas as pd
split_pictures = pd.read_csv('./validation_data/validation/validation_data.csv')

# -- Code Cell --
import os

# -- Code Cell --
os.listdir("./")

# -- Code Cell --
int(img_combined[1,1,0])

# -- Code Cell --
np.ndarray([int(img_combined[1,1,0]),int(img_combined[1,1,1]), int(img_combined[1,1,2])])

# -- Code Cell --
tile1 = cv.imread(f"./validation_data/validation/validation_tiles/105878.jpg")
tile2 = cv.imread(f"./validation_data/validation/validation_tiles/144060.jpg")
img_combined = np.zeros((128, 256, 3), dtype=np.uint8)
img_combined[0:128, 0:128] = tile1
img_combined[0:128, 128:256] = tile2
img_rgb = cv.cvtColor(img_combined, cv.COLOR_BGR2RGB)
plt.imshow(img_rgb)
plt.show()

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

# -- Code Cell --
img = cv.imread("./training_data/train_img_01.jpg")
import matplotlib.pyplot as plt
clipped_img = []
for i in range(3):
    for j in range(3):
        roi = img[128*i:128*(i+1), 128*j:128*(j+1)]
        clipped_img.append(roi)
for i in range(len(clipped_img)):
    plt.imshow(clipped_img[i])
    plt.show()

# -- Code Cell --
class ValidationDataset(Dataset):
    def __init__(self, transform, path=None):
        self.transform = transform
        if path:
            self.df = pd.read_csv(path)
    def __len__(self):
        return len(self.df)
    def __getitem__(self,index):
        row = self.df.iloc[index]
        tile1 = cv.imread(f"./validation_data/validation/validation_tiles/{row['tile1']}.jpg")
        tile2 = cv.imread(f"./validation_data/validation/validation_tiles/{row['tile2']}.jpg")
        
        img_combined = np.zeros((128, 256, 3), dtype=np.uint8)
        img_combined[0:128, 0:128] = tile1
        img_combined[0:128, 128:256] = tile2
        img_rgb = cv.cvtColor(img_combined, cv.COLOR_BGR2RGB)
        img_pil = Image.fromarray(img_rgb)
        label = row['answer']
        if self.transform:
            img_rgb = self.transform(img_pil)
        return img_rgb,label

# -- Code Cell --
valid_ds = ValidationDataset(path="./validation_data/validation/validation_data.csv", transform=transform)
valid_loader =DataLoader(valid_ds, batch_size=32, num_workers=0, shuffle=True)

# -- Code Cell --
weight = models.ResNet50_Weights.IMAGENET1K_V1

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

# -- Code Cell --
model = models.resnet50(weights=weight)
model.fc = nn.Linear(2048, 6)
for param in model.parameters():
    param.requires_grad = True
for param in model.fc.parameters():
    param.requires_grad = True
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr = 1e-4)
criterion = nn.CrossEntropyLoss()
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=15, gamma=0.5)

# -- Code Cell --
for epoch in range(25):
    model.train()       
    total_loss = 0
    correct = 0
    total = 0
 
    for combined, label in valid_loader:
        combined, label = combined.to(device), label.to(device)  
        out = model(combined)                  
        loss = criterion(out, label)        

        optimizer.zero_grad()           
        loss.backward()                
        optimizer.step()                
 
        total_loss += loss.item()
        correct += (out.argmax(1) == label).sum().item() 
        total += label.size(0)                             
    scheduler.step()  
    
    print(f"Epoc {epoch+1} -- Loss {total_loss/len(valid_loader)}")

# -- Code Cell --
class TestDataset(Dataset):
    def __init__(self,transform,path):
        self.transform = transform
        self.df = pd.read_csv(path)
    def __len__(self):
        return len(self.df)
    def __getitem__(self,index):
        row = self.df.iloc[index]
        tile1 = cv.imread(f"./test/test_tiles/{row['tile1']}.jpg")
        tile2 = cv.imread(f"./test/test_tiles/{row['tile2']}.jpg")
        img_combined = np.zeros((128, 256, 3), dtype=np.uint8)
        img_combined[0:128, 0:128] = tile1
        img_combined[0:128, 128:256] = tile2
        img_rgb = cv.cvtColor(img_combined, cv.COLOR_BGR2RGB)
        img_pil = Image.fromarray(img_rgb)
        if self.transform:
            img_rgb = self.transform(img_pil)
        return img_rgb

# -- Code Cell --
test_ds = TestDataset(path="./test/test_pairs.csv", transform=transform)
test_loader =DataLoader(test_ds, batch_size=32, num_workers=0, shuffle=False)

# -- Code Cell --
preds= []
for img in test_loader:
    with torch.no_grad():
        img = img.to(device)
        pred = model(img)
        pred = torch.argmax(pred,1)
        preds.append(pred.cpu())
        preds_final = torch.cat(preds)
preds_final = preds_final.numpy()
preds_final

# -- Code Cell --
test_pairs = pd.read_csv("./test/test_pairs.csv")

# -- Code Cell --
subs1_adica_tot_2 = pd.DataFrame({
    'subtaskID':1,
    'datapointID':test_pairs['datapointID'],
    'answer':preds_final
})
subs2 = pd.DataFrame({
    'subtaskID':2,
    'datapointID':test_pairs['datapointID'],
    'answer':preds_final
})
sub_final = pd.concat([subs1_adica_tot_2,subs2])
sub_final = sub_final.to_csv('subs.csv',index=False)