# -- Code Cell --
import pandas as pd
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 torch.nn.functional as F

# -- Code Cell --
train = pd.read_csv('./train.csv')

# -- Code Cell --
train.head()

# -- Code Cell --
import os

# -- Code Cell --
type(train['Label'][1])

# -- Code Cell --
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(os.path.join("./",f"{train['ImagePath'][index]}")).convert('RGB')
        label = row['Label']
        if self.transform:
            img = self.transform(img)
        return img, label

# -- Code Cell --
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])

# -- Code Cell --
test_transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])

# -- Code Cell --
from torch.utils.data import random_split
from sklearn.metrics import f1_score

# -- Code Cell --
train_ds = TrainDataset("./train.csv", transform=train_transform)
train_loader =DataLoader(train_ds, batch_size=32, num_workers=0, shuffle=True)

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

# -- Code Cell --
weights = models.ResNet50_Weights.IMAGENET1K_V2
model = models.resnet50(weights=weights)
model.fc = nn.Linear(2048,2)
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()

# -- Code Cell --
for epoch in range(10):
    total_loss = 0
    for img,label in train_loader:
        img = img.to(device)
        label = label.to(device)
        optimizer.zero_grad()
        output = model(img)
        loss = criterion(output,label)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"EPOCH{epoch+1}, LOSS{total_loss/len(train_loader)}")

# -- Code Cell --
test = pd.read_csv('test.csv')

# -- Code Cell --
class TestDatset(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):
        img = Image.open(os.path.join("./",f"{test['ImagePath'][index]}")).convert('RGB')
        if self.transform:
            img = self.transform(img)
        return img

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

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

# -- Code Cell --
preds_final

# -- Code Cell --
test.head()

# -- Code Cell --
subs = pd.DataFrame({
    'SampleID': test['SampleID'],
    'Label':preds_final
}).to_csv('subs.csv',index=False)

# -- Code Cell --
