###### ------------------- zkušební skript DP -----------------

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

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

# ### pomocné proměnné
# load_model = Config.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)

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

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

partition = {"testing": np.concatenate((af_names, non_af_names))}
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):')
threshold_std = Config.threshold_std
threshold_max = Config.threshold_max
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)


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


