###### ------------------- funkce pro uložení a načtení parametrů modelu DP -----------------

## Importy
import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torchvision.datasets as datasets
import torchvision.transforms as transforms

"References: https://github.com/aladdinpersson/Machine-Learning-Collection/blob/804c45e83b27c59defb12f0ea5117de30fe25289/ML/Pytorch/Basics/pytorch_loadsave.py#L26-L34"

def save_checkpoint(state, filename="my_checkpoint.pth.tar"):
    "Uložení parametrů modelu"

    print("=> Saving checkpoint")
    torch.save(state, filename)


def load_checkpoint(checkpoint, model, optimizer, scheduler):
    "Načtení parametrů modelu"

    print("=> Loading checkpoint")
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
    epoch = checkpoint['epoch']
    train_acc = checkpoint['train_acc']
    train_loss = checkpoint['train_loss']
    train_f1 = checkpoint['train_f1']
    valid_acc = checkpoint['valid_acc']
    valid_loss = checkpoint['valid_loss']
    valid_f1 = checkpoint['valid_f1']
    lrs = checkpoint['lrs']
    train_conf_matrix = checkpoint['train_conf_matrix']
    valid_conf_matrix = checkpoint['valid_conf_matrix']
    return model, optimizer, scheduler, epoch, train_acc, train_loss, train_f1, valid_acc, valid_loss, valid_f1, lrs, train_conf_matrix, valid_conf_matrix
