# 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 matplotlib.pyplot as plt

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, predict_patient_sig

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

# Path to DATA and MODEL
INPUT_DATA_PATH = r"data/data_QTDB/INFERENCE"
PATH_TO_MODEL_WEIGHTS = r"trained_models/LSTM_EKG_Model_inference.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

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

fvz = 250

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

# Initialize seed (deterministic)
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 dataset for one patient signal and his loader
patient_dataset = EKG_dataset(data_path=INPUT_DATA_PATH,
                              window_size=window_size,
                              window_overlap=window_overlap,
                              label_dict=label_dict,
                              shuffle=False)
patient_loader = DataLoader(dataset=patient_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, 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

# ###################################################
# ######### Patient signal - Evaluation #############
# ###################################################

# Compute accuracy, F1, recall and precision for one patient signal
metrics = evaluate_model_metrics(model=net, dataset_loader=patient_loader, label_dict=label_dict, device=device)

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("############################################\n"
      "################ RESULTS ###################\n"
      "############################################")

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))

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

# Save signal prediction
patient_sig_dict = predict_patient_sig(model=net,
                                       dataset_loader=patient_loader,
                                       label_dict=label_dict,
                                       device=device,
                                       overlap=window_overlap)
df_patient_sig = pd.DataFrame(patient_sig_dict)
df_patient_sig.to_csv(r"results/patient_signals.csv")

# ###################################################
# ################ Visualization ####################
# ###################################################
#
# # For predictions
# predictions_img = plt.figure()
#
# for batch_idx, (values, labels) in enumerate(patient_loader):
#     values = values.to(device)
#
#     probabilities = net(values)
#
#     predictions = torch.argmax(probabilities, axis=1)
#     predictions = torch.squeeze(torch.squeeze(predictions))
#
#     predictions_P = np.zeros(len(predictions))
#     predictions_QRS = np.zeros(len(predictions))
#     predictions_T = np.zeros(len(predictions))
#     predictions_n = np.zeros(len(predictions))
#
#     predictions_P = predictions == 0
#     predictions_QRS = predictions == 1
#     predictions_T = predictions == 2
#
#     predictions_P = predictions_P.double() * window_overlap
#     predictions_QRS = predictions_QRS.double() * window_overlap
#     predictions_T = predictions_T.double() * window_overlap
#
#     signal = values[0][0]
#     signal = signal - (sum(signal)/len(signal))
#     time_line = np.arange(
#                           batch_idx * (window_size - window_overlap),
#                           len(signal) * (batch_idx+1) - batch_idx * window_overlap
#                           ) / fvz
#     plt.plot(time_line, signal, color='b')
#
#     if batch_idx > 0:
#         plt.plot(time_line, predictions_P, color='g', label="P vlna")
#         plt.plot(time_line, predictions_QRS, color='r', label="QRS vlna")
#         plt.plot(time_line, predictions_T, color='y', label="T vlna")
#         plt.plot(time_line, predictions_n, color='k')
#         break
#     else:
#         plt.plot(time_line, predictions_P, color='g')
#         plt.plot(time_line, predictions_QRS, color='r')
#         plt.plot(time_line, predictions_T, color='y')
#         plt.plot(time_line, predictions_n, color='k')
#
#
# plt.legend(loc='upper left', bbox_to_anchor=(1.05, 1))
# plt.ylabel("[-]")
# plt.xlabel("čas [s]")
# plt.suptitle("Ukázka signálu s predikovanými třídami")
# plt.show()
# predictions_img.savefig(r"results/predictions_img.png", bbox_inches='tight')
#
#
# # For label
# label_img = plt.figure()
# for batch_idx, (values, labels) in enumerate(patient_loader):
#     values = values.to(device)
#     labels = labels.to(device)
#     labels = torch.squeeze(torch.squeeze(labels))
#
#     labels_P = np.zeros(len(labels))
#     labels_QRS = np.zeros(len(labels))
#     labels_T = np.zeros(len(labels))
#     labels_n = np.zeros(len(labels))
#
#     labels_P = labels == 0
#     labels_QRS = labels == 1
#     labels_T = labels == 2
#
#     labels_P = labels_P.double() * window_overlap
#     labels_QRS = labels_QRS.double() * window_overlap
#     labels_T = labels_T.double() * window_overlap
#
#     signal = values[0][0]
#     signal = signal - (sum(signal) / len(signal))
#     time_line = np.arange(
#                           batch_idx * (window_size - window_overlap),
#                           len(signal) * (batch_idx + 1) - batch_idx * window_overlap
#                           ) / fvz
#     plt.plot(time_line, signal, color='b')
#
#     if batch_idx > 0:
#         plt.plot(time_line, labels_P, color='g', label="P vlna")
#         plt.plot(time_line, labels_QRS, color='r', label="QRS vlna")
#         plt.plot(time_line, labels_T, color='y', label="T vlna")
#         plt.plot(time_line, labels_n, color='k')
#         break
#     else:
#         plt.plot(time_line, labels_P, color='g')
#         plt.plot(time_line, labels_QRS, color='r')
#         plt.plot(time_line, labels_T, color='y')
#         plt.plot(time_line, labels_n, color='k')
#
# plt.legend(loc='upper left', bbox_to_anchor=(1.05, 1))
# plt.ylabel("[-]")
# plt.xlabel("čas [s]")
# plt.suptitle("Ukázka signálu s anotací")
# plt.show()
# label_img.savefig(r"results/label_img.png", bbox_inches='tight')
