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

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

# %%
import os

# %%
os.listdir("./")

# %%
int(img_combined[1,1,0])

# %%
np.ndarray([int(img_combined[1,1,0]),int(img_combined[1,1,1]), int(img_combined[1,1,2])])

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

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

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

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

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

# %%
weight = models.ResNet50_Weights.IMAGENET1K_V1

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

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

# %%
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)}")

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

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

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

# %%
test_pairs = pd.read_csv("./test/test_pairs.csv")

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


