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

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

import os
import copy
import random
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import time

import torch
import torch.nn as nn
from torch import optim
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

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

# Start time measurement
tic = time.perf_counter()

# Initialize paths to TRAIN and VALIDATION datasets, and directory to save trained models
# TRAIN_DATA_PATH = r"data/data_QTDB/TRAIN"
# VALID_DATA_PATH = r"data/data_QTDB/VALID"
# SAVE_WEIGHTS_PATH = r"trained_models/"

TRAIN_DATA_PATH = r"data/data_FNKV/TRAIN"
VALID_DATA_PATH = r"data/data_FNKV/VALID"
SAVE_WEIGHTS_PATH = r"trained_models/"

# Model configuration initialization
device_string = None
num_workers = 0
batch_size = 400  # 200 400 600
epoch_count = 60
optimizer_type = "Adam"  # "Adam", "AdamW" "ASGD"
learning_rate = 0.001  # 0.001 0.005 0.01
learning_decay = 0  # 0 0.001 0.01
window_size = 150  # 300 500 700
window_overlap = 50  # 25 50 100

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

VAL_EVAL_EPOCHS = 5  # validation period epoch number (number of epochs after which, is computed validation loss)

# Weight for weighting loss values for each class individually

# For P, QRS, T, n classification
# class_weights = torch.Tensor([1.0, 1.0, 1.0, 1.0])
# class_weights = torch.Tensor([1/0.1, 1/0.15, 1/0.25, 1/0.5])
# class_weights = torch.Tensor([1/0.06, 1/0.5, 1/0.15, 1/0.29])
# class_weights = torch.Tensor([0.06, 0.5, 0.15, 0.29])

# Label dictionary to set class labels
# label_dict = {"P": 0,
#               "QRS": 1,
#               "T": 2,
#               "n": 3}

# For QRSon, QRSoff classification
class_weights = torch.Tensor([1.0, 1.0])
# class_weights = torch.Tensor([1/0.25, 1/0.75])
# class_weights = torch.Tensor([1/0.7, 1/0.3])
# class_weights = torch.Tensor([0.7, 0.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)

# Create and load TRAIN and VALIDATION datasets, into loaders
train_dataset = EKG_dataset(data_path=TRAIN_DATA_PATH,
                            window_size=window_size,
                            window_overlap=window_overlap,
                            label_dict=label_dict,
                            shuffle=True)
train_loader = DataLoader(dataset=train_dataset,
                          batch_size=batch_size,
                          num_workers=num_workers)
valid_dataset = EKG_dataset(data_path=VALID_DATA_PATH,
                            window_size=window_size,
                            window_overlap=window_overlap,
                            label_dict=label_dict,
                            shuffle=False)
valid_loader = DataLoader(dataset=valid_dataset,
                          batch_size=batch_size,
                          num_workers=num_workers)

# Enable gradient
torch.set_grad_enabled(True)

# Negative log likelihood loss function
class_weights.to(device)
loss_func = nn.NLLLoss(weight=class_weights)

# ###################################################
# ############# Model initialization ################
# ###################################################

# Setup model from architecture configuration
out_dims = len(label_dict.keys())
net = Ekg_SegNET(act_func=act_func, out_dims=out_dims, lstm_layers=lstm_layers, fc_layers=fc_layers)

optimizer_dict = {"Adam": optim.Adam, "AdamW": optim.AdamW, "ASGD": optim.ASGD}
optimizer_constructor = optimizer_dict[optimizer_type]
optimizer = optimizer_constructor(net.parameters(), lr=learning_rate, weight_decay=learning_decay)
net.to(device)

# Initialize variables for measuring loss progress
old_val_loss = evaluate_model(model=net,
                              dataset_loader=valid_loader,
                              device=device,
                              loss_func=loss_func)  # Initialize variable for measuring validation loss
progress_val_loss = [old_val_loss]  # Initialize list for logging validation loss
progress_train_loss = []  # Initialize list for logging training loss
epoch_best = None  # Initialize variable for marking best model on VALID dataset training epoch

# Initialize variable for best model weights
best_weights = copy.deepcopy(net.state_dict())  # initialize best weights variable

# ##################################################
# ################ Train Model #####################
# ##################################################

# Epoch loop
for epoch in range(epoch_count):

    net.train()  # Check to enable grad, after validation

    # Print current epoch
    print("EPOCH - " + str(epoch + 1) + "/" + str(epoch_count))

    # Batch loop
    for batch_idx, (values, labels) in enumerate(train_loader):

        # Set inputs to processing unit
        print("TRAIN BATCH: " + str(batch_idx))
        values = values.to(device)
        labels = labels.to(device)

        # Propagate input and compute loss from output and labels
        with torch.cuda.amp.autocast():  # Solves unknown CUDA problem
            probabilities = net(values)
            loss = loss_func(torch.log(probabilities), labels)

        # Print current average batch loss
        print("CURRENT BATCH LOSS: " + str(loss.item()))

        # Backpropagate weigth update
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print("EPOCH - " + str(epoch + 1) + "/" + str(epoch_count))
        print("PROGRESS OF EPOCH " + str(batch_idx) + "/" + str(np.ceil(len(train_dataset) / batch_size)))

        if len(progress_train_loss) == 0:
            progress_train_loss.append(loss)

    # Validate model on VALIDATION dataset each VAL_EVAL_EPOCHS epochs
    if (epoch + 1) % VAL_EVAL_EPOCHS == 0:
        print("Validating model")
        val_loss = evaluate_model(model=net, dataset_loader=valid_loader, device=device, loss_func=loss_func)

        # Save current validation adn training loss progress
        progress_val_loss.append(val_loss)
        progress_train_loss.append(loss)

        print("Old validation loss: " + str(old_val_loss))
        print("New validation loss: " + str(val_loss))

        # Save model in current state
        torch.save(net.state_dict(), SAVE_WEIGHTS_PATH + "/LSTM_EKG_Model_weights" + str(epoch) + ".pt")

        # Rewrite old validation loss and save best model
        if val_loss < old_val_loss:
            old_val_loss = val_loss
            epoch_best = epoch
            best_weights = copy.deepcopy(net.state_dict())
        else:
            old_val_loss = val_loss

# Save last and best model weights
torch.save(net.state_dict(), SAVE_WEIGHTS_PATH + "/LSTM_EKG_Model_weights_LAST" + str(epoch_count) + ".pt")
torch.save(best_weights, SAVE_WEIGHTS_PATH + "/LSTM_EKG_Model_weights_BEST" + str(epoch_best) + ".pt")

toc = time.perf_counter()
print("Time: ")
print((toc - tic), "\n")

# Plot loss progress and save plot to .png
x_axis = list(range(0, epoch_count + 1, VAL_EVAL_EPOCHS))

progress_val_loss = [float(x) for x in progress_val_loss]
progress_train_loss = [float(x.detach().cpu()) for x in progress_train_loss]

print(x_axis)
print(progress_val_loss)
print(progress_train_loss)

df_loss = pd.DataFrame({"Epoch": x_axis, "Val_loss": progress_val_loss, "Train_loss": progress_train_loss})
df_loss.to_csv(r"results/train_val_loss.csv")

# loss_img = plt.figure()
# plt.plot(x_axis, progresfs_val_loss)
# plt.ylabel('Výstup ztrátová funkce')
# plt.xlabel('epochy učení')
# plt.suptitle('Vývoj ztrátové funkce')
# plt.show()
# loss_img.savefig(r"results/Validation loss.png", bbox_inches='tight')
