# !pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

# !pip install scikit-image

import copy
import os
import pickle
import time
import warnings
from functools import partial
from multiprocessing.pool import ThreadPool
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torchvision
from matplotlib import colors
from matplotlib import pyplot as plt
from PIL import Image
from skimage import io
from sklearn.preprocessing import LabelEncoder
from torch.nn import functional as F
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, models, transforms
from torchvision.io import read_image
from tqdm import tqdm, tqdm_notebook

%matplotlib inline

warnings.filterwarnings(action="ignore", category=DeprecationWarning)

from sklearn.model_selection import train_test_split

train_val_labels = [path.parent.name for path in train_val_files]
train_files, val_files = train_test_split(train_val_files, test_size=0.20, stratify=train_val_labels, random_state=SEED)
train_labels_raw, val_labels_raw = [path.parent.name for path in train_files], [path.parent.name for path in val_files]

label_encoder = LabelEncoder()
label_encoder.fit(train_labels_raw)

train_labels = label_encoder.transform(train_labels_raw)
val_labels = label_encoder.transform(val_labels_raw)
test_labels = None

assert len(train_labels) == len(train_files)

class2label = dict(zip(label_encoder.classes_, label_encoder.transform(label_encoder.classes_)))
label2class = {v: k for k, v in class2label.items()}

val_test_trainsform = transforms.Compose(
    [
        transforms.ToPILImage(),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)
train_transform = transforms.Compose(
    [
        transforms.ToPILImage(),
        transforms.Resize((224, 224)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)

data_transforms = {
    "train": train_transform,
    "val": val_test_trainsform,
    "test": val_test_trainsform,
}

class SimpsonsDataset(Dataset):
    def __init__(self, file_names, img_labels, mode, transform=None, target_transform=None):
        self.file_names = file_names
        self.img_labels = img_labels
        self.transform = transform
        self.target_transform = target_transform
        self.mode = mode

    def __len__(self):
        return len(self.file_names)

    def __getitem__(self, idx):
        img_path = str(self.file_names[idx])
        image = read_image(img_path)

        if self.transform:
            image = self.transform(image)

        if self.mode == "test":
            return image
        else:
            label = self.img_labels[idx]

            if self.target_transform:
                label = self.target_transform(label)

            return image, label

train_dataset = SimpsonsDataset(train_files, train_labels, "train", data_transforms["train"])
val_dataset = SimpsonsDataset(val_files, val_labels, "val", data_transforms["val"])
test_dataset = SimpsonsDataset(test_files, test_labels, "test", data_transforms["test"])

image_datasets = {"train": train_dataset, "val": val_dataset}
dataset_sizes = {x: len(image_datasets[x]) for x in ["train", "val"]}

dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)

batch_size = 64
epochs = 10

train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

dataloaders = {"train": train_dataloader, "val": val_dataloader}

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print(f"Epoch {epoch}/{num_epochs - 1}")
        print("-" * 10)

        # Each epoch has a training and validation phase
        for phase in ["train", "val"]:
            if phase == "train":
                model.train()  # Set model to training mode
            else:
                model.eval()  # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in tqdm_notebook(dataloaders[phase]):
                inputs = inputs.to(device)
                labels = labels.type(torch.LongTensor).to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == "train"):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == "train":
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            if phase == "train":
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print(f"{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}")

            # deep copy the model
            if phase == "val" and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time() - since
    print(f"Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s")
    print(f"Best val Acc: {best_acc:4f}")

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model

model_conv = torchvision.models.vgg16_bn(pretrained=True)
for param in model_conv.parameters():
    param.requires_grad = False

num_classes = len(class2label)
model_conv.classifier = nn.Linear(512 * 7 * 7, num_classes)
model_conv = model_conv.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model_conv.parameters(), lr=1e-3, amsgrad=True)
exp_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)

def predict(model, test_loader):
    with torch.no_grad():
        logits = []

        for inputs in tqdm_notebook(test_loader):
            inputs = inputs.to(device)
            model.eval()
            outputs = model(inputs).cpu()
            logits.append(outputs)

    probs = nn.functional.softmax(torch.cat(logits), dim=-1).numpy()
    return probs

test_filenames = [path.name for path in test_dataset.file_names]

submission = pd.read_csv("data/sample_submission.csv")
submission = pd.DataFrame({"Id": test_filenames, "Expected": preds}).sort_values("Id")
submission.to_csv("./submission.csv", index=False)