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

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

import os
import numpy as np
import pandas as pd
import random

import torch
from torch.utils.data import DataLoader

from src.architecture.model import Ekg_SegNET
from src.data_loading.EKG_Dataset import EKG_dataset
from src.utils.utilities import evaluate_model_metrics

# ###################################################
# ################ Initialization ###################
# ###################################################

# Path to TEST DATA and MODEL to test
# TEST_DATA_PATH = r"data/data_QTDB/TEST"
# TEST_DATA_PATH = r"data/data_QTDB/TRAIN"
TEST_DATA_PATH = r"data/data_QTDB/VALID"
PATH_TO_MODEL_WEIGHTS = r"trained_models/LSTM_EKG_Model.pt"

# Model configuration initialization
device_string = None
num_workers = 0
batch_size = 400  # 200 400 600
window_size = 300  # 300 500 700
window_overlap = 25   # 25 50 100
out_dims = 2

act_func = "relu"  # "gelu"
lstm_layers = 1  # 2 3
fc_layers = 1  # 2

# Class label initialization
label_dict = {"P": 0,
              "QRS": 1,
              "T": 2,
              "n": 3}

# label_dict = {"QRS": 1,
#               "n": 0}

# Initialize seed
seed = 42

os.environ['PYTHONHASHSEED'] = '0'
np.random.seed(seed)
random.seed(seed)
torch.random.manual_seed(seed)
torch.cuda.manual_seed(seed=seed)
torch.backends.cudnn.deterministic = True

# Pick device to compute on
if not device_string:
    device_string = "cuda" if torch.cuda.is_available() else "cpu"

device = torch.device(device_string)

# Initialize TEST dataset and loader
test_dataset = EKG_dataset(data_path=TEST_DATA_PATH,
                           window_size=window_size,
                           window_overlap=window_overlap,
                           label_dict=label_dict,
                           shuffle=False)
test_loader = DataLoader(dataset=test_dataset,
                         batch_size=batch_size,
                         num_workers=num_workers)

# Initialize model and load weights from path PATH_TO_MODEL_WEIGHTS
net = Ekg_SegNET(act_func=act_func, out_dims=out_dims, lstm_layers=lstm_layers, fc_layers=fc_layers)
net.to(device)
net.load_state_dict(torch.load(PATH_TO_MODEL_WEIGHTS, map_location=device))
net.eval()  # turn off Batch Normalization layers and gradient


# ####################################################
# ############## Evaluate TEST dataset ###############
# ####################################################

# Compute accuracy, F1, recall and precision of TEST data set
metrics = evaluate_model_metrics(model=net, dataset_loader=test_loader, label_dict=label_dict, device=device)

# Save results into csv
df_loss = pd.DataFrame(metrics)
df_loss.to_csv(r"results/TEST_DATASET_evaluation_metrics.csv")

average_acc = np.average(list(metrics["accuracy"].values()))
average_recall = np.average(list(metrics["recall"].values()))
average_precision = np.average(list(metrics["precision"].values()))
average_F1 = np.average(list(metrics["F1"].values()))

print("TEST dataset accuracy is: " + str(average_acc))
print("TEST dataset recall is: " + str(average_recall))
print("TEST dataset precision is: " + str(average_precision))
print("TEST dataset F1 is: " + str(average_F1))
