# -*- coding: utf-8 -*-
"""diplomka_neuronka_kod.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1r102t9--PhLz1MZWE7gMGviKZAlFdIeK
"""

import torch
import PIL
import time
import copy
from PIL import Image
import numpy as np
import torchvision
from torch import nn, optim, cuda
from torch.optim import lr_scheduler
from torchvision import transforms, models, datasets
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from collections import Counter
import pandas as pd
from timeit import default_timer as timer
import os
import tqdm
import itertools


from google.colab import files
from google.colab import drive

"""Načítanie prepravených dát z My Drive, augmentácia a pripravenie dataLoaderov pre trénovanie/validáciu/testovanie"""

drive.mount('/content/gdrive')
# path = /content/gdrive/MyDrive/DIPLOMKA/data/dataset300

train_dir = '/content/gdrive/MyDrive/DIPLOMKA/data/datasetALL/train'
val_dir = '/content/gdrive/MyDrive/DIPLOMKA/data/datasetALL/validate'
test_dir = '/content/gdrive/MyDrive/DIPLOMKA/data/datasetALL/test'
test_neuronka_dir = '/content/gdrive/MyDrive/DIPLOMKA/data/test_neuronka'

dirs = {'train' : train_dir,
        'validate' : val_dir,
        'test' : test_dir,
        'test_neuronka' : test_neuronka_dir}


data_transforms = {
    'train' :  transforms.Compose([transforms.RandomRotation(30),
                                   transforms.RandomResizedCrop(224),
                                   transforms.ToTensor(),
                                   transforms.Normalize([0.485, 0.456, 0.406],
                                                        [0.229, 0.224, 0.225])
                                   ]),
                   
    'validate' :  transforms.Compose([transforms.Resize(256),
                                      transforms.CenterCrop(224),
                                      transforms.ToTensor(),
                                      transforms.Normalize([0.485, 0.456, 0.406],
                                                           [0.229, 0.224, 0.225])
                    ]),
    'test' :  transforms.Compose([transforms.Resize(256),
                                  transforms.CenterCrop(224),
                                  transforms.ToTensor(),
                                  transforms.Normalize([0.485, 0.456, 0.406],
                                                       [0.229, 0.224, 0.225])
                    ]),
    'test_neuronka' : transforms.Compose([transforms.Resize(256),
                                          transforms.CenterCrop(224),
                                          transforms.ToTensor(),
                                          transforms.Normalize([0.485, 0.456, 0.406],
                                                               [0.229, 0.224, 0.225])
                    ])

}

image_dataset = {x: datasets.ImageFolder(dirs[x], transform =  data_transforms[x])
for x in ['train', 'validate', 'test', 'test_neuronka']}

dataloaders = {x: torch.utils.data.DataLoader(image_dataset[x], batch_size=32, shuffle=True) for x in ['train', 'validate', 'test', 'test_neuronka']}
dataset_sizes = {x: len(image_dataset[x]) for x in ['train', 'validate', 'test', 'test_neuronka']}

#dataloaders1 = {x: torch.utils.data.DataLoader(image_dataset[x], batch_size=16, shuffle=True) for x in ['train',  'test', 'test_neuronka']}
#dataloaders2 = {x: torch.utils.data.DataLoader(image_dataset[x], batch_size=16*2, shuffle=True) for x in ['validate']}
#dataloaders.update(dataloaders1)
#dataloaders.update(dataloaders2)


class_names = image_dataset['train'].classes
print(class_names)

"""Display datasets statistics"""

labels_as_nums = image_dataset['train'].class_to_idx
num_as_label = {
    index: label
    for label, index in image_dataset['train'].class_to_idx.items()
}


train_nums = Counter([num_as_label[x] for x in image_dataset['train'].targets])
val_nums = Counter([num_as_label[x] for x in image_dataset['validate'].targets])
test_nums = Counter([num_as_label[x] for x in image_dataset['test'].targets])
test_general_nums = Counter([num_as_label[x] for x in image_dataset['test_neuronka'].targets])

train_nums = pd.DataFrame({'class' :list(train_nums.keys()), 'train_num': list(train_nums.values())})
val_nums = pd.DataFrame({'class' :list(val_nums.keys()), 'valid_num': list(val_nums.values())})
test_nums = pd.DataFrame({'class' :list(test_nums.keys()), 'test_num': list(test_nums.values())})
test_general_nums = pd.DataFrame({'class' :list(test_general_nums.keys()), 'test_different_dataset': list(test_general_nums.values())})
dataset_table = pd.merge(train_nums,val_nums,on='class',how='left').merge(test_nums,on='class',how='left').merge(test_general_nums,on='class',how='left')

dataset_table.head()

"""Nastavenie trénovania modelu na GPU"""

train_on_gpu = cuda.is_available()
print(f'Train on gpu: {train_on_gpu}')
multi_gpu = False
# Number of gpus
if train_on_gpu:
    gpu_count = cuda.device_count()
    print(f'{gpu_count} gpus detected.')
    if gpu_count > 1:
        multi_gpu = True

print(train_on_gpu,multi_gpu)

"""Train model, save model and plot results functions

Vráti natrénovaný model a dáta na vykreslenie trénovacej/validačnej presnosti a chyby

Uloží model na Drive


"""

def train_model(model,
          criterion,
          optimizer,
          train_loader,
          valid_loader,
          save_file_name,
          max_epochs_stop,
          n_epochs,
          print_every):


    # počet epoch, kedy nenastale žiadne zlepšenie modelu, predčasné ukončenie trénovania
    epochs_no_improve = 0
    valid_loss_min = np.Inf

    valid_max_acc = 0
    history = []


    #počet epoch, na koľkých už bol model trénovaný
    try:
        print(f'Model has been trained for: {model.epochs} epochs.\n')
    except:
        model.epochs = 0
        print(f'Starting Training from Scratch.\n')

    overall_start = timer()

    
    
    for epoch in range(n_epochs):

        # štatistika počas trénovania
        train_loss = 0.0
        valid_loss = 0.0

        train_acc = 0
        valid_acc = 0

        # Začiatok trénovania
        model.train()
        start = timer()

       
        for ii, (data, target) in enumerate(train_loader):
            
            if train_on_gpu:
                data, target = data.cuda(), target.cuda()

            optimizer.zero_grad()
            output = model(data)

            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            train_loss += loss.item() * data.size(0)

            # výpočet presnosti
            _, pred = torch.max(output, dim=1)
            correct_tensor = pred.eq(target.data.view_as(pred))
            accuracy = torch.mean(correct_tensor.type(torch.FloatTensor))
            train_acc += accuracy.item() * data.size(0)

            print(
                f'Epoch: {epoch}\t{100 * (ii + 1) / len(train_loader):.2f}% complete. {timer() - start:.2f} seconds elapsed in epoch.',
                end='\r')

        # začiatok validácie modelu
        else:
            model.epochs += 1

            with torch.no_grad():
                model.eval()

                for data, target in valid_loader:

                    if train_on_gpu:
                        data, target = data.cuda(), target.cuda()

                    output = model(data)

                    # výpoćet validačnej straty
                    loss = criterion(output, target)
                    valid_loss += loss.item() * data.size(0)

                    _, pred = torch.max(output, dim=1)
                    correct_tensor = pred.eq(target.data.view_as(pred))
                    accuracy = torch.mean(correct_tensor.type(torch.FloatTensor))
                    valid_acc += accuracy.item() * data.size(0)

                # Výpočet priemerných strát a presností
                train_loss = train_loss / len(train_loader.dataset)
                valid_loss = valid_loss / len(valid_loader.dataset)

                train_acc = train_acc / len(train_loader.dataset)
                valid_acc = valid_acc / len(valid_loader.dataset)

                history.append([train_loss, valid_loss, train_acc, valid_acc])

                if (epoch + 1) % print_every == 0:
                    print(
                        f'\nEpoch: {epoch} \tTraining Loss: {train_loss:.4f} \tValidation Loss: {valid_loss:.4f}'
                    )
                    print(
                        f'\t\tTraining Accuracy: {100 * train_acc:.2f}%\t Validation Accuracy: {100 * valid_acc:.2f}%'
                    )

                # Uloženie modelu pri poklese validačnej straty
                if valid_loss < valid_loss_min:
                    torch.save(model.state_dict(), save_file_name)
                    epochs_no_improve = 0
                    valid_loss_min = valid_loss
                    valid_best_acc = valid_acc
                    best_epoch = epoch

                # Inkrementácia počtu epoch bez zlepšenia
                else:
                    epochs_no_improve += 1
                    if epochs_no_improve >= max_epochs_stop:
                        print(
                            f'\nEarly Stopping! Total epochs: {epoch}. Best epoch: {best_epoch} with loss: {valid_loss_min:.2f} and acc: {100 * valid_best_acc:.2f}%'
                        )
                        total_time = timer() - overall_start
                        print(
                            f'{total_time:.2f} total seconds elapsed. {total_time / (epoch+1):.2f} seconds per epoch.'
                        )

                        model.load_state_dict(torch.load(save_file_name))
                        model.optimizer = optimizer

                        history = pd.DataFrame(
                            history,
                            columns=[
                                'train_loss', 'valid_loss', 'train_acc',
                                'valid_acc'
                            ])
                        return model, history


    model.optimizer = optimizer
    total_time = timer() - overall_start
    print(
        f'\nBest epoch: {best_epoch} with loss: {valid_loss_min:.2f} and acc: {100 * valid_best_acc:.2f}%'
    )
    print(
        f'{total_time:.2f} total seconds elapsed. {total_time / (epoch+1):.2f} seconds per epoch.'
    )
    history = pd.DataFrame(
        history,
        columns=['train_loss', 'valid_loss', 'train_acc', 'valid_acc'])
    return model, history


#uloženie natrénovaného modelu
def save_checkpoint(model, path):


    model_name = path.split('-')[0]
    #print(model_name)
    assert (model_name in ['vgg16', 'resnet50', 'resnet18', 'resnet152'
                           ]), "Path must have the correct model name"

    
    checkpoint = {
        'class_to_idx': model.class_to_idx,
        'num_as_label': model.num_as_label,
        'epochs': model.epochs,
    }

    if model_name == 'resnet50' or model_name == 'resnet152' or model_name == 'resnet18':
        if multi_gpu:
            checkpoint['fc'] = model.module.fc
            checkpoint['state_dict'] = model.module.state_dict()
        else:
            checkpoint['fc'] = model.fc
            checkpoint['state_dict'] = model.state_dict()

    checkpoint['optimizer'] = model.optimizer
    checkpoint['optimizer_state_dict'] = model.optimizer.state_dict()

    path1 = '/content/gdrive/MyDrive/DIPLOMKA/data/'
    torch.save(checkpoint, path1+path)

# vykreslenie strát
def plot_Losses():
  figureTVLosses = plt.figure(figsize=(8, 6))
  for c in ['train_loss', 'valid_loss']:
      plt.plot(
          history[c], label=c)
  plt.legend()
  plt.xlabel('Epoch')
  plt.ylabel('Average Negative Log Likelihood')
  plt.title('Training and Validation Losses')

  file_name_TV_losses = save_file_name + 'Training and Validation Losses'

  figureTVLosses.savefig(file_name_TV_losses +'.jpg')
  #/content/Training and Validation Losses.jpg
  files.download(file_name_TV_losses + ".jpg")

# vykreslenie presností
def plot_Accuracies():
  figureTVAccs = plt.figure(figsize=(8, 6))
  for c in ['train_acc', 'valid_acc']:
      plt.plot(
          100 * history[c], label=c)
  plt.legend()
  plt.xlabel('Epoch')
  plt.ylabel('Average Accuracy')
  plt.title('Training and Validation Accuracy')

  file_name_TV_accs = save_file_name + 'Training and Validation Accuracies'

  figureTVAccs.savefig(file_name_TV_accs +'.jpg')
  #/content/Training and Validation Losses.jpg
  files.download(file_name_TV_accs + ".jpg")

# vykreslenie confusion matice
def plot_confusion_matrix():
  import seaborn as sns
  nb_classes = 5
  confusion_matrix = np.zeros((nb_classes, nb_classes))
  with torch.no_grad():
      for i, (inputs, classes) in enumerate(dataloaders['test']):
          inputs = inputs.to(device)
          classes = classes.to(device)
          outputs = my_model(inputs)
          _, preds = torch.max(outputs, 1)
          for t, p in zip(classes.view(-1), preds.view(-1)):
                  confusion_matrix[t.long(), p.long()] += 1

  confusion_matrix_figure = plt.figure(figsize=(15,10))

  #class_names = class_names
  df_cm = pd.DataFrame(confusion_matrix, index=class_names, columns=class_names).astype(int)
  heatmap = sns.heatmap(df_cm, annot=True, fmt="d")

  heatmap.yaxis.set_ticklabels(heatmap.yaxis.get_ticklabels(), rotation=0, ha='right',fontsize=15)
  heatmap.xaxis.set_ticklabels(heatmap.xaxis.get_ticklabels(), rotation=45, ha='right',fontsize=15)
  plt.ylabel('True label')
  plt.xlabel('Predicted label')

  file_name_confusion_matrix = save_file_name + 'Confusion matrix'

  confusion_matrix_figure.savefig(file_name_confusion_matrix +'.jpg')
  files.download(file_name_confusion_matrix + ".jpg")

"""Príprava modelu, výber modelu, zamrazenie vrstiev a návrh klasifikátora"""

my_model = models.resnet50(pretrained=True)

num_ftrs = my_model.fc.in_features

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

#my_model.classifier[6].out_features = 5


'''
ct = 0
for child in my_model.children():
    ct += 1
    if ct < 47:
      for param in child.parameters():
          param.requires_grad = False

print(my_model)


for param in my_model.parameters():
          param.requires_grad = False
for param in my_model.fc.parameters():
            param.require_grad = True 
'''
for param in my_model.parameters():
          param.requires_grad = False


'''
my_model.classifier = nn.Sequential(
            nn.Linear(25088, 4096),
            nn.ReLU(True),
            nn.Dropout(0.25),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(0.25),
            nn.Linear(4096, 5),
            #nn.LogSoftmax(dim=1),
        )
#print(my_model)
'''
my_model.fc = nn.Sequential(
              nn.Linear(num_ftrs, 256),
              nn.ReLU(),
              nn.Dropout(0.4),
              nn.Linear(256, 5),
              #nn.Dropout(0.4),
              nn.LogSoftmax(dim=1))


#my_model.fc = nn.Linear(num_ftrs, 5)
my_model = my_model.to(device)

my_model.class_to_idx = image_dataset['train'].class_to_idx
my_model.num_as_label = {
    idx: class_
    for class_, idx in my_model.class_to_idx.items()
}

print(my_model)


'''
my_model.fc = nn.Sequential(
    nn.Linear(num_ftrs, 256), 
    nn.ReLU(), 
    nn.Linear(256, 5)
)
'''

"""Spustenie trénovania a nastavenie cesty na uloženie výsledku"""

save_file_name = f'resnet50-46Pretrained_datasetALL_24.4.pt'
checkpoint_path = f'resnet50-46Pretrained_datasetALL_24.4.pth'

criterion = nn.CrossEntropyLoss()

optimizer_ft = optim.SGD(my_model.parameters(), lr=0.01, momentum=0.9)
#optimizer_ft = optim.Adam(my_model.parameters(), lr=1e-3, weight_decay=1e-4)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

my_model, history = train_model(my_model, criterion, optimizer_ft, dataloaders['train'], dataloaders['validate'], save_file_name=save_file_name, max_epochs_stop=35, n_epochs=100, print_every=1)
#torch.save(my_model.state_dict(), 'resnet50-Pretrained_datasetALL_6.5.pth')
#save_checkpoint(my_model, path=checkpoint_path)

class_names = image_dataset['train'].classes
print(class_names)
plot_Losses()
plot_Accuracies()

save_checkpoint(my_model, path=checkpoint_path)

"""Načítanie uloženého modelu"""

def load_checkpoint(path):
    
    if torch.cuda.is_available():
      checkpoint = torch.load(path)
    else:
      checkpoint = torch.load(path, map_location=torch.device('cpu'))
    model = models.resnet50(pretrained=True)

    for param in model.parameters():
        param.requires_grad = False
    model.fc = checkpoint['fc']

    # Load in the state dict
    model.load_state_dict(checkpoint['state_dict'])

    total_params = sum(p.numel() for p in model.parameters())
    print(f'{total_params:,} total parameters.')
    total_trainable_params = sum(
        p.numel() for p in model.parameters() if p.requires_grad)
    print(f'{total_trainable_params:,} total gradient parameters.')

    
    model = model.to(device)


   # model.class_to_idx = checkpoint['class_to_idx']
    model.num_as_label = checkpoint['num_as_label']
    #model.idx_to_class = checkpoint['idx_to_class']
    model.epochs = checkpoint['epochs']

    optimizer = checkpoint['optimizer']
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

    return model, optimizer


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
root_path = '/content/gdrive/MyDrive/DIPLOMKA/data/'
save_file_name = f'resnet50-Pretrained_datasetALL_21.4.pt'
checkpoint_path_load = f'resnet50-Pretrained_datasetALL_21.4.pth'
#/content/gdrive/MyDrive/DIPLOMKA/data/modelsresnet152-dataset1000.pth
my_model, optimizer = load_checkpoint(path=root_path+checkpoint_path_load)

print(my_model)

"""Vykreslenie finálnej štatistiky"""

def process_image(image_path):
    """Process an image path into a PyTorch tensor"""

    image = Image.open(image_path)
    img = image.resize((256, 256))

    width = 256
    height = 256
    new_width = 224
    new_height = 224

    left = (width - new_width) / 2
    top = (height - new_height) / 2
    right = (width + new_width) / 2
    bottom = (height + new_height) / 2
    img = img.crop((left, top, right, bottom))

    img = np.array(img).transpose((2, 0, 1)) / 256

    means = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))
    stds = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))

    img = img - means
    img = img / stds

    img_tensor = torch.Tensor(img)

    return img_tensor

def predict(image_path, model, topk=5):

    real_class = image_path.split('/')[-2]

    img_tensor = process_image(image_path)

    if train_on_gpu:
        img_tensor = img_tensor.view(1, 3, 224, 224).cuda()
    else:
        img_tensor = img_tensor.view(1, 3, 224, 224)

    with torch.no_grad():
        model.eval()
        out = model(img_tensor)
        ps = torch.exp(out)

        topk, topclass = ps.topk(topk, dim=1)

        top_classes = [
            model.num_as_label[class_] for class_ in topclass.cpu().numpy()[0]
        ]
        top_p = topk.cpu().numpy()[0]

        return img_tensor.cpu().squeeze(), top_p, top_classes, real_class


def imshow_tensor(image, ax=None, title=None):


    if ax is None:
        fig, ax = plt.subplots()

    image = image.numpy().transpose((1, 2, 0))

    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = std * image + mean

    image = np.clip(image, 0, 1)

    ax.imshow(image)
    plt.axis('off')

    return ax, image


np.random.seed = 100


def random_test_image(aTest_Dir):

    c = np.random.choice(class_names)
    root = aTest_Dir +'/'+ np.random.choice(class_names)+'/'
    img_path = root + np.random.choice(os.listdir(root))
    return img_path


def display_prediction(image_path, model, topk):

    img, ps, classes, y_obs = predict(image_path, model, topk)

    result = pd.DataFrame({'p': ps}, index=classes)

    plt.figure(figsize=(16, 5))
    ax = plt.subplot(1, 2, 1)
    ax, img = imshow_tensor(img, ax=ax)

    ax.set_title(y_obs, size=20)

    ax = plt.subplot(1, 2, 2)
    result.sort_values('p')['p'].plot.barh(color='blue', edgecolor='k', ax=ax)
    plt.xlabel('Predicted Probability')
    plt.tight_layout()

def accuracy(output, target, topk=(1, )):

    if train_on_gpu:
        output = output.to('cuda')
        target = target.to('cuda')

    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(k=maxk, dim=1, largest=True, sorted=True)
        pred = pred.t()

        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []

        for k in topk:
            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size).item())
        return res


def evaluate(model, test_loader, criterion, topk=(1, 2, 3)):

    confusion_matrix = torch.zeros(5, 5)
    classes = []
    losses = []

    acc_results = np.zeros((len(test_loader.dataset), len(topk)))
    i = 0

    model.eval()
    with torch.no_grad():

        for data, targets in test_loader:

            if train_on_gpu:
                data, targets = data.to('cuda'), targets.to('cuda')

            out = model(data)
            _, predscm = torch.max(out, 1)
            for t, p in zip(targets.view(-1), predscm.view(-1)):
                confusion_matrix[t.long(), p.long()] += 1
            
            for pred, true in zip(out, targets):

                acc_results[i, :] = accuracy(
                    pred.unsqueeze(0), true.unsqueeze(0), topk)
                classes.append(model.num_as_label[true.item()])

                loss = criterion(pred.view(1, 5), true.view(1))
                losses.append(loss.item())
                i += 1

    results = pd.DataFrame(acc_results, columns=[f'top{i}' for i in topk])
    results['class'] = classes
    results['loss'] = losses
    results = results.groupby(classes).mean()

    return results.reset_index().rename(columns={'index': 'class'}),confusion_matrix

criterion = nn.CrossEntropyLoss()
results,confusion_matrix = evaluate(my_model, dataloaders['test'], criterion)


def plot_confusion_matrix(cm, classes, title='Confusion matrix', cmap=plt.cm.Blues):

    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=45)
    plt.yticks(tick_marks, classes)

    fmt = '.0f' 
    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')

    file_name_confusion_matrix = save_file_name + 'Confusion matrix'
    
    #plt.savefig(file_name_confusion_matrix +'.jpg')
    #files.download(file_name_confusion_matrix + ".jpg")
    return plt

classnames = [my_model.num_as_label[i] for i in range(0,5)]
plt.figure(figsize=(15,10))
plt = plot_confusion_matrix(confusion_matrix, classnames)


results = results.merge(dataset_table,left_on='class',right_on='class')
results

for i in range(30):
  display_prediction(random_test_image(test_neuronka_dir), my_model, topk=5)