# -- Code Cell --
from PIL import Image
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset,random_split
from torchvision import models, transforms
import torch.optim as optim

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

# -- Code Cell --
validation_data

# -- Code Cell --
tile1 = Image.open(f"./validation_data/validation/validation_tiles/{validation_data['tile1'][1]}.jpg")
tile2 = Image.open(f"./validation_data/validation/validation_tiles/{validation_data['tile2'][1]}.jpg")
combined = np.zeros((128, 256, 3), dtype=np.uint8)
print(combined.shape)
combined[0:128, 0:128] = tile1
combined[0:128, 128:256] = tile2
img = Image.fromarray(combined)
import matplotlib.pyplot as plt
plt.imshow(img)

# -- Code Cell --
class TrainDataset(Dataset):
    def __init__(self, df, transform):
        self.df = df
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        tile1 = Image.open(f"./validation_data/validation/validation_tiles/{row['tile1']}.jpg")
        tile2 = Image.open(f"./validation_data/validation/validation_tiles/{row['tile2']}.jpg")
        label = row['answer']
        combined = np.zeros((128, 256, 3), dtype=np.uint8)
        combined[0:128, 0:128] = tile1
        combined[0:128, 128:256] = tile2
        img = Image.fromarray(combined)
        img=self.transform(img)
        return img, label
        

# -- Code Cell --
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])
])

# -- Code Cell --
train_ds = TrainDataset(validation_data, transform)
train_subset, val_subset = random_split(train_ds, [0.8,0.2])
train_loader =DataLoader(train_subset, batch_size=32, shuffle=True,num_workers=0)
val_loader =DataLoader(val_subset, batch_size=32, shuffle=True,num_workers=0)

# -- Code Cell --
validation_data['answer'].nunique()

# -- Code Cell --
class conv(nn.Module):
    def __init__(self,in_c, out_c):
        super().__init__()
        self.c1 = nn.Conv2d(in_c, out_c,kernel_size=3, stride=1 )
        self.bn = nn.BatchNorm2d(out_c)
        self.relu = nn.ReLU(inplace=True)
    def forward(self,x):
        c1 = self.c1(x)
        bn = self.bn(c1)
        out = self.relu(bn)
        return out

# -- Code Cell --
class modelmeu(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.c1 = conv(3,32)
        self.c2 = conv(32, 64)
        self.c3 = conv(64,128)
        self.c4 = conv(128, 256)
        self.pool = nn.MaxPool2d(2)
        self.out = nn.Sequential(
            nn.Flatten(),
            nn.LazyLinear(6)
        )
    def forward(self,x):
        x = self.pool(self.c1(x))
        x = self.pool(self.c2(x))
        x = self.pool(self.c3(x))
        x = self.pool(self.c4(x))
        return self.out(x)

# -- Code Cell --
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
model.fc = nn.LazyLinear(6)
for param in model.parameters():
    param.requires_grad = True
optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max')
device ='cuda'
model=model.to(device)

# -- Code Cell --
from sklearn.metrics import f1_score

# -- Code Cell --
decoy = 10
for epoch in range(20):
    model.train()
    for img, label in train_loader:
        img,label = img.to(device), label.to(device)
        optimizer.zero_grad()
        output = model(img)
        loss = criterion(output,label)
        loss.backward()
        optimizer.step()
    model.eval()
    all_labels = []
    all_preds = []
    for img, label in val_loader:
        img,label = img.to(device), label.to(device)
        output = model(img)
        pred = output.argmax(dim=1)
        all_preds.append(pred.cpu())
        all_labels.append(label.cpu())
    all_preds = torch.cat(all_preds).numpy()
    all_labels = torch.cat(all_labels).numpy()
    f1 = f1_score(all_labels,all_preds, average='macro')
    if f1>decoy:
        decoy = f1
        torch.save(model.state_dict(), "model.pth")
    scheduler.step(f1)
    print(f"f1  {f1}")

# -- Code Cell --
model.load_state_dict(torch.load("model.pth"))
model.eval()

# -- 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 = Image.open(f"./test/test_tiles/{row['tile1']}.jpg")
        tile2 = Image.open(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_pil = Image.fromarray(img_combined)
        if self.transform:
            img_pil = self.transform(img_pil)
        return img_pil

# -- 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 --
sub1 = pd.DataFrame({
    'subtaskID':1,
    'datapointID':test['datapointID'],
    'answer':answer
})
sub2 = pd.DataFrame({
    'subtaskID':2,
    'datapointID':test['datapointID'],
    'answer':answer
})
final = pd.concat([sub1,sub2]).to_csv("subs.csv", index=False)