###### ------------------- trénovací, validační a testovací smyčka modelu DP -----------------

### importy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import f1_score, classification_report, confusion_matrix, precision_recall_curve

from filereader import load_file


### inicializace kriteriální funkce
def cross_entropy(res, lbls):

    return -torch.mean(torch.sum((lbls * torch.log(res)), dim=1))

### trénování a validace
def train_valid(model, loader, device, optimizer, input):

    "trénování a validace za 1 epochu; input = 'training'/'validation'/'testing'"

    if input == "training":
        ## trénovací cyklus
        model.train() ## Set model to train mode
        loss_temp, acc_temp, f1_temp, conf_matrix_temp = main_cycle(model, loader, device, optimizer, input)
    elif input == "validation":
        ## validační cyklus
        with torch.no_grad():
            model.eval()  ## ovlivňuje BN
            loss_temp, acc_temp, f1_temp, conf_matrix_temp = main_cycle(model, loader, device, optimizer, input)
    elif input == "testing":
        ## testovací cyklus
        with torch.no_grad():
            model.eval()  ## ovlivňuje BN
            loss, acc, f1, clas_temp, true_label, heat_maps, conf_matrix = main_cycle(model, loader, device, optimizer, input)

    if input == "testing":
        return loss, acc, f1, clas_temp, true_label, heat_maps, conf_matrix
    else:
        ## průměrná chyba a přesnost
        epoch_loss = np.mean(loss_temp)  ## chyba za celou epochu na trénovací/validační množině dat
        epoch_acc = np.mean(acc_temp)  ## přesnost za celou epochu na trénovací/validační množině dat
        epoch_f1 = np.mean(f1_temp)  ## f1 skore za celou epochu na trénovací/validační množině dat

        return epoch_loss, epoch_acc, epoch_f1, conf_matrix_temp



def main_cycle(model, loader, device, optimizer, input):

    ## pomocné proměnné pro výpočet přesnosti a chyby
    loss_temp = list()
    acc_temp = list()
    f1_temp = list()
    conf_matrix_temp = list()

    for index_t, (x_batch, y_batch, lengths) in enumerate(loader):
        ## vložení dat na grafickou kartu, pokud je přítomna
        x_batch = x_batch.to(device=device)
        y_batch = y_batch.to(device=device)
        lengths = lengths.to(device=device)
        ## výstup sítě = Forward pass - Predicted outputs
        predictions, heatmap = model(x_batch, lengths)

        ## zobrazení a ověření parametrů
        # for name, param in model.named_parameters():
        #     if param.requires_grad:
        #         print(name, param.data)

        ## výpočet chyby - kriteriální funkce
        loss = cross_entropy(predictions, y_batch)

        if input == "training":
            ## Clear gradients
            optimizer.zero_grad()
            ## zpětná propagace = back-propagation --> výpočet gradientů
            loss.backward()
            ## Update the parameters (weights)
            optimizer.step()

        ##  -------------------- pro 2 třídy, varianta se softmax --------------------
        y_batch_num = np.argmax(y_batch.detach().cpu().numpy(), axis=1)
        clas = np.argmax(predictions.detach().cpu().numpy(), axis=1)

        ## výpočet přesnosti
        acc = np.mean((clas == y_batch_num)).astype(np.float32)

        ## výpočet matice záměn a F1 skóre
        conf_matrix, f1 = classification_metrics(clas, y_batch_num)

        ## ---------------

        loss_temp.append(loss.detach().cpu().numpy())
        acc_temp.append(acc)
        f1_temp.append(f1)
        conf_matrix_temp.append(conf_matrix)

        if input == "testing":
            loss = loss.detach().cpu().numpy()
            heat_maps = heatmap.detach().cpu().numpy()

    if input == "testing":
        return loss, acc, f1, clas, y_batch_num, heat_maps, conf_matrix
    else:
        return loss_temp, acc_temp, f1_temp, conf_matrix_temp


def classification_metrics(classification, ground_truth):

    """
    Funkce pro výpočet matice záměn, F1 skóre popř. precision (PPH) a recall (senzitivita).
    """
    ## výpočet přesnosti
    # acc = np.mean((classification == ground_truth)).astype(np.float32)

    ## výpočet confusion matrix (matice záměn)
    TP, FP, TN, FN = 0, 0, 0, 0
    if np.size(classification) > 1:
        for y, y_clas in zip(ground_truth, classification):
            if y == 1 and y_clas == 1:
                TP += 1
            elif y == 0 and y_clas == 0:
                TN += 1
            elif y == 1 and y_clas == 0:
                FN += 1
            elif y == 0 and y_clas == 1:
                FP += 1
    else:
        if ground_truth == 1 and classification == 1:
            TP += 1
        elif ground_truth == 0 and classification == 0:
            TN += 1
        elif ground_truth == 1 and classification == 0:
            FN += 1
        elif ground_truth == 0 and classification == 1:
            FP += 1

    conf_matrix = np.array([[TN, FP],
                            [FN, TP]])

    ## výpočet f1 skore
    if (TP + FP) == 0:
        precision = 0
    else:
        precision = TP / (TP + FP)
    if (TP + FN) == 0:
        recall = 0
    else:
        recall = TP / (TP + FN)
    if (precision + recall) == 0:
        f1 = 0
    else:
        f1 = np.float32(2 * (precision * recall) / (precision + recall))

    return conf_matrix, f1
