"""
## Setup
This notebook requires some packages besides pytorch-lightning.
"""
# !pip install --quiet "torchvision" "torch" "pytorch-lightning" "torchmetrics"
# !pip install --quiet comet-ml
import comet_ml

from kaggle_secrets import UserSecretsClient
from pytorch_lightning.loggers import CometLogger
user_secrets = UserSecretsClient()
comet_logger = CometLogger(
    api_key=user_secrets.get_secret("COMET_API_KEY"),
    workspace=user_secrets.get_secret("COMET_WORKSPACE"),  # Optional
    save_dir=".",  # Optional
    project_name=user_secrets.get_secret("COMET_PROJECT"),  # Optional
    # rest_api_key=user_secrets.get_secret("COMET_REST_API_KEY"),  # Optional
    # experiment_name="smap",  # Optional
)
import os

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchmetrics
from PIL import Image
from torch.utils.data import DataLoader
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from torchvision import datasets
from torchvision import models
from torchvision import transforms
from pytorch_lightning import Trainer
import pandas as pd
from pytorch_lightning.utilities.seed import seed_everything
from sklearn.model_selection import train_test_split
ROOT = os.path.join("/kaggle", "input", "lego-minifigures-classification")
AVAIL_GPUS = min(1, torch.cuda.device_count())
BATCH_SIZE = 4
seed_everything(42)
"""
## Simplest example

Here's the simplest most minimal example with just a training loop (no validation, no testing).

**Keep in Mind** - A `LightningModule` *is* a PyTorch `nn.Module` - it just has a few more helpful features.
"""
class Model(pl.LightningModule):
    def __init__(self, classes: int):
        super(Model, self).__init__()
        
        self.classes = classes
        self.create_model()

        self.train_acc = torchmetrics.Accuracy()
        self.valid_acc = torchmetrics.Accuracy()
        self.test_acc = torchmetrics.Accuracy()

    def create_model(self):
#        self.models = models.resnet50(pretrained=True)
        self.models = models.mobilenet_v2(pretrained=True)
    
        print(self.models)

        # num_ftrs = self.models.fc.in_features
        # self.models.fc = nn.Linear(num_ftrs, self.classes)
        num_ftrs = self.models.classifier[1].in_features
        self.models.classifier[1] = nn.Linear(num_ftrs, self.classes)

    def forward(self, x):
        return self.models(x)

    def training_step(self, batch, batch_idx):
        images, target = batch
        preds = self(images)

        loss = F.cross_entropy(preds, target)
        
        _, preds = torch.max(preds, 1)
        self.train_acc(preds, target)
        self.log('train_acc', self.train_acc, on_step=True, on_epoch=True, prog_bar=True)

        return loss

    def validation_step(self, batch, batch_idx):
        images, target = batch
        preds = self.forward(images)
        
        loss = F.cross_entropy(preds, target)
        
        _, preds = torch.max(preds, 1)
        self.valid_acc(preds, target)
        self.log('valid_acc', self.valid_acc, on_step=False, on_epoch=True, prog_bar=True)
        return loss

    def test_step(self, batch, batch_idx):
        images, target = batch
        preds = self.forward(images)
        
        loss = F.cross_entropy(preds, target)
        
        _, preds = torch.max(preds, 1)
        self.test_acc(preds, target)
        self.log('test_acc', self.test_acc)
        self.log("test_loss", loss)
        return loss

    def configure_optimizers(self):
        # return torch.optim.SGD(self.parameters(), lr=0.01, momentum=0.9)
        return torch.optim.Adam(self.parameters(), lr=0.0001)
data = pd.read_csv(os.path.join(ROOT,'index.csv'))
data.head()
class LegoDataset(torch.utils.data.Dataset):
    def __init__(self, X: list, y: list, root: str, transforms = None):
        self.X = list(X)
        self.y = list(y)

        self.root = root
        self.transforms = transforms

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

    def __getitem__(self, index):
        filename = os.path.join(self.root, self.X[index])
        label = self.y[index]
        image = Image.open(filename)
        
        if self.transforms is not None:
            image = self.transforms(image)
        
        return image, label
X, y = data.path, data.class_id
y = y-1

X_train, X_valid, y_train, y_valid = train_test_split(X, y, random_state=0)
print(len(X_train), len(X_valid))
train_dataset = LegoDataset(X_train, y_train, root=ROOT, transforms=transforms.Compose([
                                    transforms.Resize(224),
                                    transforms.ToTensor(),
                                    transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
                                   ])
                           )
                            
train_dataset[0]
valid_dataset = LegoDataset(X_valid, y_valid, root=ROOT, transforms=transforms.Compose([
                                    transforms.Resize(224),
                                    transforms.ToTensor(),
                                    transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
                                   ])
                           )

valid_dataset[0]
train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE)
valid_dataloader = DataLoader(valid_dataset, batch_size=1)
classes = data.class_id.unique()
print(classes, len(classes))
model = Model(classes=max(classes))
#AVAIL_GPUS=0

trainer = Trainer(
    logger=comet_logger,
    callbacks=[EarlyStopping(monitor="valid_acc", mode="max", patience=5, min_delta=0.00)],
    gpus=AVAIL_GPUS,
    max_epochs=50,
    default_root_dir="./test/"
)

hist = trainer.fit(model, train_dataloader, valid_dataloader)
test_data = pd.read_csv(os.path.join(ROOT, 'test.csv'))
valid_dataset = LegoDataset(test_data.path, test_data.class_id-1, root=ROOT, transforms=transforms.Compose([
                                    transforms.Resize(224),
                                    transforms.ToTensor(),
                                    transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
                                   ])
                           )

test_dataloader = DataLoader(valid_dataset, batch_size=1)
trainer.test(model, test_dataloader)
"""
# Evaluate
"""
import itertools
import numpy as np
import matplotlib.pyplot as plt

def plot_confusion_matrix(cm, classes, normalize=False, title='Confusion matrix', cmap=plt.cm.Blues):
    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        print("Normalized confusion matrix")
    else:
        print('Confusion matrix, without normalization')
    fig, ax = plt.subplots(figsize=(len(classes), len(classes)))
    print(cm)
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=90)
    plt.yticks(tick_marks, classes)

    fmt = '.2f' if normalize else 'd'
    thresh = cm.max() / 2.
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        plt.text(j, i, format(cm[i, j], fmt), horizontalalignment="center", color="white" if cm[i, j] > thresh else "black")

    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
import torchmetrics
device = torch.device("cuda")

model.eval()
model = model.cuda()

class_names = pd.read_csv(os.path.join(ROOT, 'metadata.csv'))["minifigure_name"].tolist()

acc = torchmetrics.Accuracy()

cm_p = []
cm_t = []

for i, (inputs, targets) in enumerate(test_dataloader):
    inputs = inputs.to(device)

    preds = model(inputs).cpu()
    _, preds = torch.max(preds, 1)
    acc(preds, targets)

    #      MC
    cm_t.extend(targets.cpu().detach().numpy())
    cm_p.extend(preds.cpu().detach().numpy())

from sklearn.metrics import confusion_matrix
cm = confusion_matrix(cm_t, cm_p)

plot_confusion_matrix(cm, class_names, normalize=False)
plot_confusion_matrix(cm, class_names, normalize=True)
comet_logger.experiment.log_confusion_matrix(cm_t, cm_p, labels=class_names)
comet_logger.experiment.end()