# DEEP LEARNING BASED QRS DELINEATOR
# Malina Ondrej Bc. 
# FEKT UBMI 2020/2021

# ###################################################
# #################### Imports ######################
# ###################################################

import numpy as np
import torch

# ###################################################
# ############# Evaluation functions ################
# ###################################################


# Evaluates model loss on data in dataset object
def evaluate_model(model, dataset_loader, device, loss_func):
    model.eval()  # Dont use batch normalization layers
    loss_vector = []

    # Without keeping gradient
    with torch.no_grad():
        # Batch loop
        for batch_idx, (values, labels) in enumerate(dataset_loader):
            # Set inputs to processing unit
            values = values.to(device)
            labels = labels.to(device)

            # Propagate input and compute loss from output and labels
            with torch.cuda.amp.autocast():
                probabilities = model(values)
                loss_vector.append(loss_func(torch.log(probabilities), labels))

        # Set to torch tensor and return loss as a sum of batch losses
        torch_loss = torch.FloatTensor(loss_vector)

        return torch.sum(torch_loss)


# Evaluates model metrics on data in dataset object
def evaluate_model_metrics(model, dataset_loader, label_dict, device):
    model.eval()

    # Preparation of a dictionary of evaluation metrics
    accuracy_dict = {}
    recall_dict = {}
    precision_dict = {}
    F1_dict = {}

    div_zero_const = 0.00000001  # Prevents division by zero
    label_strings = label_dict.keys()

    TP = {}
    for key in label_strings:
        TP[key] = 0

    TN = TP.copy()
    FP = TP.copy()
    FN = TP.copy()

    # Calculation of the confusion matrix and evaluation metrics for each class
    with torch.no_grad():
        for batch_idx, (values, labels) in enumerate(dataset_loader):

            print("TEST BATCH " + str(batch_idx) + "/" + str(int(np.ceil(len(dataset_loader)))))

            values = values.to(device)
            labels = labels.to(device)

            labels = torch.squeeze(torch.squeeze(labels))
            probabilities = model(values)

            predictions = torch.argmax(probabilities, axis=1)
            predictions = torch.squeeze(torch.squeeze(predictions))

            for class_str in label_strings:
                class_label = label_dict[class_str]

                label_true = labels == class_label
                label_false = labels != class_label
                pred_true = predictions == class_label
                pred_false = predictions != class_label

                TP[class_str] += sum(sum(label_true * pred_true)).cpu().numpy()
                TN[class_str] += sum(sum(label_false * pred_false)).cpu().numpy()
                FP[class_str] += sum(sum(label_false * pred_true)).cpu().numpy()
                FN[class_str] += sum(sum(label_true * pred_false)).cpu().numpy()

    for class_str in label_strings:
        accuracy_dict[class_str] = (TP[class_str] + TN[class_str]) / (
                TP[class_str] + TN[class_str] +
                FP[class_str] + FN[class_str] +
                div_zero_const
        )
        recall_dict[class_str] = TP[class_str] / (TP[class_str] + FN[class_str] + div_zero_const)
        precision_dict[class_str] = TP[class_str] / (TP[class_str] + FP[class_str] + div_zero_const)
        F1_dict[class_str] = 2 / (
                (1 / (precision_dict[class_str] + div_zero_const))
                +
                (1 / (recall_dict[class_str] + div_zero_const))
                + div_zero_const
        )

    metrics = {
        "accuracy": accuracy_dict,
        "recall": recall_dict,
        "precision": precision_dict,
        "F1": F1_dict
    }

    return metrics


# Evaluates patient signal on model and saves prediction signal
def predict_patient_sig(model, dataset_loader, device, overlap):
    model.eval()

    label_sig = []
    pred_sig = []

    # Calculation of the confusion matrix and evaluation metrics for each class
    with torch.no_grad():
        for batch_idx, (values, labels) in enumerate(dataset_loader):
            print("PATIENT BATCH " + str(batch_idx) + "/" + str(int(np.ceil(len(dataset_loader)))))

            values = values.to(device)
            labels = labels.to(device)

            labels = torch.squeeze(torch.squeeze(labels))
            probabilities = model(values)

            predictions = torch.argmax(probabilities, axis=1)
            predictions = torch.squeeze(torch.squeeze(predictions))

            without_overlap = labels.shape[1] - overlap

            labels_cuts = torch.split(labels, [without_overlap, overlap], dim=1)
            predictions_cuts = torch.split(predictions, [without_overlap, overlap], dim=1)
            label_sig.extend(torch.flatten(labels_cuts[0].detach().cpu()).numpy())
            pred_sig.extend(torch.flatten(predictions_cuts[0].detach().cpu()).numpy())

        label_sig.extend(labels_cuts[1].detach().cpu().numpy()[-1])
        pred_sig.extend(predictions_cuts[1].detach().cpu().numpy()[-1])
    return {"Predicted signal": pred_sig, "Ground truth": label_sig}
