###### ------------------- hlavní soubor DP -----------------

"Reference: https://github.com/aladdinpersson/Machine-Learning-Collection/blob/157a5f458f272a513eb6b4a19d6613aec32dc21c/ML/Pytorch/Basics/pytorch_simple_CNN.py#L25-L41"

### importy

import numpy as np
import random
import matplotlib.pyplot as plt
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision.transforms as transforms
import time
from torch.utils.data import DataLoader, Dataset
from bayes_opt import BayesianOptimization
from scipy.signal import medfilt, find_peaks, peak_prominences, peak_widths
from scipy import interpolate
from sklearn.metrics import mean_squared_error

from model import ResNet
from config import Config
from datareader import data_reader, collate_fn
from data_preprocessing import data_preprocessing, data_distribution
from train_valid import train_valid, classification_metrics
from loadsave import load_checkpoint, save_checkpoint
from afib_classification import parox_afib_classification
from afib_localization import create_feature_signals, parox_afib_localization, evaluation_of_localization
from plot_maps import plot_heat_maps

###########################################################################################################
###### ------------------------ Klasifikace na FiS a Ostatní pomocí ResNet ------------------------- ######
###########################################################################################################

### ------------ nastavení zařízení ------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

### -------------- hyperparametry --------------
batch_size = Config.batch_size
learning_rate = Config.learning_rate
num_epochs = Config.num_epochs

### -------------- parametry modelu --------------
levels = Config.levels                  ## počet residuálních bloků (= ResNet) = 6
lvl1_size = Config.lvl1_size            ## počet filtrů konvolučních vrstev uvnitř 1. residuálního bloku - o násobky se zvětšuje hloubka v blocích = 12
convs_in_layer = Config.convs_in_layer  ## počet konvolučních vrstev uvnitř jednotlivých bloků = 4
init_conv = Config.init_conv            ## počet filtrů v úvodní konvoluční vrstvě - init_conv = výstupní počet příznaků pro self.init_conv (1. konv. vrstva) = 12
filter_size = Config.filter_size        ## velikost kernelu = 3

### -------------- načtení a přezpracování dat --------------
path = Config.path_train
file_of_sample_names = Config.file_of_sample_names
ecg_data_dict, labels_dict, labels_name_dict, signal_of_af_location_dict, qrs_loc_dict, loc_af_start_dict, loc_af_end_dict, fs, af_names, non_af_names = data_preprocessing(path, file_of_sample_names)

### -------------- rozdělení dat --------------
partition = data_distribution(af_names, non_af_names, split_index_train=0.8, split_index_test=0.1)

### -------------- načtení trénovacích a validačních dat --------------
training_dataset = data_reader("training", partition, ecg_data_dict, labels_dict)
training_loader = DataLoader(training_dataset, batch_size=batch_size, shuffle=True, drop_last=True, collate_fn=collate_fn)

validation_dataset = data_reader("validation", partition, ecg_data_dict, labels_dict)
validation_loader = DataLoader(validation_dataset, batch_size=batch_size, shuffle=True, drop_last=False, collate_fn=collate_fn)

### -------------- inicializace modelu --------------
model = ResNet(levels=levels, lvl1_size=lvl1_size, input_size=Config.input_size,
               output_size=Config.output_size, convs_in_layer=convs_in_layer, init_conv=init_conv,
               filter_size=filter_size, mil_solution=Config.mil_solution)

### -------------- inicializace optimalizačního algoritmus a scheduler pro změnu kroku učení --------------
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)

###################### --- Trénovací cyklus --- ############################

### pomocné proměnné
load_model = Config.load_model

## načtení parametrů trénování
if load_model:
    ## načtení parametrů trénování
    load_model_name = Config.load_model_name
    model, optimizer, scheduler, load_epoch, train_acc, train_loss, train_f1, valid_acc, valid_loss, valid_f1, lrs, train_conf_matrix, valid_conf_matrix = \
        load_checkpoint(torch.load(load_model_name, map_location=device), model, optimizer, scheduler)
    start_epoch = load_epoch
    model = model.to(device)
    print(Config.load_model_name)
else:
    train_loss = list()
    train_acc = list()
    train_f1 = list()
    valid_loss = list()
    valid_acc = list()
    valid_f1 = list()
    lrs = list()
    train_conf_matrix = list()
    valid_conf_matrix = list()
    start_epoch = 0
    model = model.to(device)
    print(Config.model_name)

## ------------ trénování neuronové sítě ---------------
## čas začátku trénování
start_time = time.time()
for epoch in range(start_epoch, num_epochs):

    ## trénovací cyklus
    train_loss_temp, train_acc_temp, train_f1_temp, train_conf_matrix_temp = \
        train_valid(model, training_loader, device, optimizer, input="training")
    ## validační cyklus
    valid_loss_temp, valid_acc_temp, valid_f1_temp, valid_conf_matrix_temp = \
        train_valid(model, validation_loader, device, optimizer, input="validation")

    lrs.append(optimizer.param_groups[0]["lr"]) ## uložení velikosti kroku učení
    scheduler.step()                            ## změna kroku učení

    ## uložení trénovací průměrné chyby, přesnosti a f1 skore do listu
    train_loss.append(train_loss_temp)  ## chyba za celou epochu na trénovací množině dat
    train_acc.append(train_acc_temp)  ## přesnost za celou epochu na trénovací množině dat
    train_f1.append(train_f1_temp)  ## f1 skore za celou epochu na trénovací množině dat

    ## uložení validační průměrné chyby, přesnosti a f1 skore do listu
    valid_loss.append(valid_loss_temp)  ## chyba za celou epochu na validační množině dat
    valid_acc.append(valid_acc_temp)  ## přesnost za celou epochu na validační množině dat
    valid_f1.append(valid_f1_temp)  ## f1 skore za celou epochu na validační množině dat

    ## uložení poslední matice záměn v epoše pro trénovací a valiační množinu
    train_conf_matrix.append(train_conf_matrix_temp[-1])
    valid_conf_matrix.append(valid_conf_matrix_temp[-1])

    ## získání parametrů modelu + trénování - uložení parametrů současné epochy
    if (epoch + 1) % 2 == 0:
        checkpoint = {'epoch': epoch + 1,
                      'model_state_dict': model.state_dict(),
                      'optimizer_state_dict': optimizer.state_dict(),
                      'scheduler_state_dict': scheduler.state_dict(),
                      'train_acc': train_acc,
                      'train_loss': train_loss,
                      'train_f1': train_f1,
                      'valid_acc': valid_acc,
                      'valid_loss': valid_loss,
                      'valid_f1': valid_f1,
                      'lrs': lrs,
                      'train_conf_matrix': train_conf_matrix,
                      'valid_conf_matrix': valid_conf_matrix}
        ## průběžné uložení parametrů modelu a trénování
        if epoch + 1 >= 46:
            save_checkpoint(checkpoint, filename=f'{Config.model_name[:-8]}_epoch_{epoch + 1}.pth.tar')
        else:
            save_checkpoint(checkpoint)

    # if (epoch + 1) % 2 == 0:
        ## vypsání výsledků epochy do konzole
        print(f'Epoch {epoch + 1}/{num_epochs} \t train_loss = {train_loss[epoch]} \t valid_loss = {valid_loss[epoch]}')
        # print(f'Epoch {epoch + 1}/{num_epochs} \t train_acc = {train_acc[epoch]} \t valid_acc = {valid_acc[epoch]}')
        print(f'Epoch {epoch + 1}/{num_epochs} \t train_f1 = {train_f1[epoch]} \t valid_f1 = {valid_f1[epoch]}')
## čas ukončení trénování
end_time = time.time()

## uložení parametrů modelu a trénování podle názvu modelu
if not load_model:
    save_checkpoint(checkpoint, filename=Config.model_name)

last_valid_f1 = np.mean(valid_f1[-5:])  # .astype(np.float32) ## průměr f1 z posledních 5 epoch
print('Mean of last 5 valid F1', last_valid_f1)

## výpočet průměrných chyb, přesností a f1 skore na celé trénovací množině
last_train_acc = train_acc[-1]
last_train_loss = train_loss[-1]
last_train_f1 = train_f1[-1]
## výpočet průměrných chyb, přesností a f1 skore na celé validační množině
last_valid_acc = valid_acc[-1]
last_valid_loss = valid_loss[-1]
last_valid_f1 = valid_f1[-1]

print(f'Last: \t train_loss = {last_train_loss} \t valid_loss = {last_valid_loss}')
print(f'Last: \t train_acc = {last_train_acc} \t valid_acc = {last_valid_acc}')
print(f'Last: \t train_f1 = {last_train_f1} \t valid_f1 = {last_valid_f1}')

print(f'Last: \t train_conf_matrix = {train_conf_matrix[-1]} \n valid_conf_matrix = {valid_conf_matrix[-1]}')


## počet učících parametrů modelu
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'Number of trainable parameters: {total_params}')

## celková doba trvání trénovacího cyklu
print('Time of training:', end_time - start_time)

## vykreslení chyby
plt.figure()
plt.plot(train_loss, label="Trénovací množina")
plt.plot(valid_loss, label="Validační množina")
plt.xlabel("Epocha [-]")
plt.ylabel("Chyba [-]")
plt.legend()
# plt.title(f"Loss-{Config.model_name[:-8]}")
plt.show()

## vykreslení přesnosti
# plt.figure()
# plt.plot(train_acc, label="Train accuracy")
# plt.plot(valid_acc, label="Valid accuracy")
# plt.title("Accuracy")
# plt.legend()
# plt.show()

## vykreslení f1 skore
plt.figure()
plt.plot(train_f1, label="Trénovací množina")
plt.plot(valid_f1, label="Validační množina")
plt.xlabel("Epocha [-]")
plt.ylabel("F1 skóre [-]")
plt.legend()
# plt.title(f"F1-{Config.model_name[:-8]}")
plt.show()

## vykreslení kroku učení
# plt.figure()
# plt.plot(lrs)
# plt.xlabel("Velikost kroku učení [-]")
# plt.ylabel("F1 skóre [-]")
# plt.title("Krok učení během trénování")
# plt.show()

###########################################################################################################
###### --------- Klasifikace parox a persi FiS a lokalizace parox FiS - trénovací množina ---------- ######
###########################################################################################################

## --- vytvoření df pro lokalizaci
training_loader_loc = DataLoader(training_dataset, batch_size=2000, shuffle=False, drop_last=False, collate_fn=collate_fn)
_, _, _, train_clas, train_true_label, train_heat_maps, _ = train_valid(model, training_loader_loc, device, optimizer, input="testing")

train_names = list()
for ind in training_dataset.list_IDs:
    train_names.append(labels_name_dict[ind]) ## get value - name of labels

## výpočet matice záměn pro klasifikaci FiS
train_conf_mat_afib, train_f1_afib = classification_metrics(train_clas, train_true_label)
print(f'Training dataset - afib classification:\n{train_conf_mat_afib}, F1 = {train_f1_afib}')

df_dict = {
            "clas_FiS":         train_clas,
            "true_label_FiS":   train_true_label,
            "names":            training_dataset.list_IDs.tolist(),
            "names_of_labels":  train_names
            }

train_summary_df = pd.DataFrame(df_dict)

train_df = train_summary_df[train_summary_df["clas_FiS"] == 1].copy()
train_df['true_label_parox'] = np.where(train_df['names_of_labels'] == 'paroxysmal atrial fibrillation', 1, 0)

## vykreslení heat map
# heat_map_train_ind = train_df[train_df["clas_FiS"] == 1][20:40].index.tolist()
# plot_heat_maps(train_heat_maps, train_df, heat_map_train_ind)

## ----------- Klasifikace paroxysmálních a perzistentní FiS - rozhodovací pravidla -----------
heat_map_train_ind = train_df[train_df["clas_FiS"] == 1].index.tolist()
ground_truth_parox_train = train_df['true_label_parox'].to_list()

print('Training dataset - Evaluation of classification (paroxysmal afib):')
## klasifikace s prohledáváním prostoru
# clas_parox_train, index_parox_train, index_non_parox_train, threshold_std, threshold_max = \
#     parox_afib_classification(train_heat_maps, heat_map_train_ind, ground_truth_parox_train, input="training")
## klasifikace bez prohledávání prostoru
threshold_std = Config.threshold_std
threshold_max = Config.threshold_max
clas_parox_train, index_parox_train, index_non_parox_train = \
    parox_afib_classification(train_heat_maps, heat_map_train_ind, ground_truth_parox_train, input="testing", threshold_std=threshold_std, threshold_max=threshold_max)


## ----------- Vytvoření feature signálu -----------
fsigs_train = create_feature_signals(train_heat_maps, index_parox_train, train_df, ecg_data_dict, fs, signal_of_af_location_dict)

## ----------- Lokalizace paroxysmálních FiS -----------
det_start_all_list_train, det_end_all_list_train, det_start_all_true_list_train, det_end_all_true_list_train = \
    parox_afib_localization(fsigs_train, index_parox_train, train_df, qrs_loc_dict, loc_af_start_dict, loc_af_end_dict, ecg_data_dict, fs, signal_of_af_location_dict)

## vyhodnocení lokalizace
print('Training dataset - Evaluation of paroxysmal localization:')
evaluation_of_localization(det_start_all_list_train, det_end_all_list_train, det_start_all_true_list_train, det_end_all_true_list_train, fs)




###########################################################################################################
###### --------- Klasifikace parox a persi FiS a lokalizace parox FiS - validační množina ---------- ######
###########################################################################################################

## --- vytvoření df pro lokalizaci
validation_loader_loc = DataLoader(validation_dataset, batch_size=2000, shuffle=False, drop_last=False, collate_fn=collate_fn)
_, _, _, valid_clas, valid_true_label, valid_heat_maps, _ = train_valid(model, validation_loader_loc, device, optimizer, input="testing")

valid_names = list()
for ind in validation_dataset.list_IDs:
    valid_names.append(labels_name_dict[ind]) ## get value - name of labels

## výpočet matice záměn pro klasifikaci FiS
valid_conf_mat_afib, valid_f1_afib = classification_metrics(valid_clas, valid_true_label)
print(f'Validation dataset - afib classification:\n{valid_conf_mat_afib}, F1 = {valid_f1_afib}')

df_dict = {
            "clas_FiS":         valid_clas,
            "true_label_FiS":   valid_true_label,
            "names":            validation_dataset.list_IDs.tolist(),
            "names_of_labels":  valid_names
            }

valid_summary_df = pd.DataFrame(df_dict)

valid_df = valid_summary_df[valid_summary_df["clas_FiS"] == 1].copy()
valid_df['true_label_parox'] = np.where(valid_df['names_of_labels'] == 'paroxysmal atrial fibrillation', 1, 0)

## vykreslení heat map
# heat_map_valid_ind = valid_df[valid_df["clas_FiS"] == 1][0:20].index.tolist()
# plot_heat_maps(valid_heat_maps, valid_df, heat_map_valid_ind)

## ----------- Klasifikace paroxysmálních a perzistentní FiS - rozhodovací pravidla -----------
heat_map_valid_ind = valid_df[valid_df["clas_FiS"] == 1].index.tolist()
ground_truth_parox_valid = valid_df['true_label_parox'].to_list()

print('Validation dataset - Evaluation of classification (paroxysmal FiS):')
clas_parox_valid, index_parox_valid, index_non_parox_valid = \
    parox_afib_classification(valid_heat_maps, heat_map_valid_ind, ground_truth_parox_valid, input="testing", threshold_std=threshold_std, threshold_max=threshold_max)

## ----------- Vytvoření feature signálu -----------
fsigs_valid = create_feature_signals(valid_heat_maps, index_parox_valid, valid_df, ecg_data_dict, fs, signal_of_af_location_dict)

## ----------- Lokalizace paroxysmálních FiS -----------
det_start_all_list_valid, det_end_all_list_valid, det_start_all_true_list_valid, det_end_all_true_list_valid = \
    parox_afib_localization(fsigs_valid, index_parox_valid, valid_df, qrs_loc_dict, loc_af_start_dict, loc_af_end_dict, ecg_data_dict, fs, signal_of_af_location_dict)

## vyhodnocení lokalizace
print('Validation dataset - Evaluation of paroxysmal localization:')
evaluation_of_localization(det_start_all_list_valid, det_end_all_list_valid, det_start_all_true_list_valid, det_end_all_true_list_valid, fs)






###########################################################################################################
###### --------- Klasifikace parox a persi FiS a lokalizace parox FiS - Testovací množina ---------- ######
###########################################################################################################

###################### --- Testování - klasifikace FiS a Ostatní --- ############################

testing_dataset = data_reader("testing", partition, ecg_data_dict, labels_dict)
testing_loader = DataLoader(testing_dataset, batch_size=1000, shuffle=False, drop_last=False, collate_fn=collate_fn)

## testovací cyklus
test_loss, test_acc, test_f1, test_clas, test_true_label, heat_maps, test_conf_matrix = train_valid(model, testing_loader, device, optimizer, input="testing")

print(f'Testing: \t test_loss = {test_loss} \t test_acc = {test_acc} \t test_f1 = {test_f1} \n test_conf_matrix = {test_conf_matrix}')

test_names = list()
for ind in testing_dataset.list_IDs:
    test_names.append(labels_name_dict[ind]) ## get value - name of labels

## výpočet matice záměn pro klasifikaci FiS
test_conf_mat_afib, test_f1_afib = classification_metrics(test_clas, test_true_label)
print(f'Testing dataset - afib classification:\n{test_conf_mat_afib}, F1 = {test_f1_afib}')

df_dict = {
            "clas_FiS": test_clas,
            "true_label_FiS": test_true_label,
            "names": testing_dataset.list_IDs.tolist(),
            "names_of_labels": test_names
            }

test_summary_df = pd.DataFrame(df_dict)

test_df = test_summary_df[test_summary_df["clas_FiS"] == 1].copy()
test_df['true_label_parox'] = np.where(test_df['names_of_labels'] == 'paroxysmal atrial fibrillation', 1, 0)

## vykreslení heat map
# heat_map_ind = test_df[test_df["clas_FiS"] == 1][0:15].index.tolist()
# plot_heat_maps(heat_maps, test_df, heat_map_ind)

########################### ----------- Segmentace paroxysmální FiS ----------- ###########################

## ----------- Klasifikace paroxysmálních a perzistentní FiS - rozhodovací pravidla -----------
heat_map_test_ind = test_df[test_df["clas_FiS"] == 1].index.tolist()
ground_truth_parox_test = test_df['true_label_parox'].to_list()

print('Testing dataset - Evaluation of classification (paroxysmal FiS):')
clas_parox_test, index_parox_test, index_non_parox_test = \
    parox_afib_classification(heat_maps, heat_map_test_ind, ground_truth_parox_test, input="testing", threshold_std=threshold_std, threshold_max=threshold_max)

## ----------- Vytvoření feature signálu -----------
fsigs_test = create_feature_signals(heat_maps, index_parox_test, test_df, ecg_data_dict, fs, signal_of_af_location_dict)

## ----------- Lokalizace paroxysmálních FiS -----------
det_start_all_list_test, det_end_all_list_test, det_start_all_true_list_test, det_end_all_true_list_test = \
    parox_afib_localization(fsigs_test, index_parox_test, test_df, qrs_loc_dict, loc_af_start_dict, loc_af_end_dict, ecg_data_dict, fs, signal_of_af_location_dict)

## vyhodnocení lokalizace
print('Testing dataset - Evaluation of paroxysmal localization:')
evaluation_of_localization(det_start_all_list_test, det_end_all_list_test, det_start_all_true_list_test, det_end_all_true_list_test, fs)


##########################################################################