##### ------------------- lokalizace fibrilace síní z feature signálů 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
from torch.utils.data import DataLoader
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 train_valid import classification_metrics
from config import Config

## ----------- Vytvoření feature signálu odstranění  -----------

def create_feature_signals(heat_maps, index_parox, test_df, ecg_data_dict, fs, signal_of_af_location_dict):

    fsigs = list()

    for i in index_parox:

        name_of_fsig = test_df.names[i]
        ## vytvoření osy x
        x = np.linspace(0, ecg_data_dict[name_of_fsig].shape[1] / fs, ecg_data_dict[name_of_fsig].shape[1])
        x_fsig = np.linspace(0, ecg_data_dict[name_of_fsig].shape[1] / fs, heat_maps[i].shape[1])

        ## volba vhodného feature signálu
        hm_max = np.max(heat_maps[i, :, :])
        hm_fsigs_max = np.max(heat_maps[i, :, :], axis=1)
        threshold_fsig = 0.75 * hm_max
        fsig_val = hm_fsigs_max[hm_fsigs_max > threshold_fsig]
        fsig_ind = np.where(hm_fsigs_max > threshold_fsig)[0]

        fsigs_concat_med = np.median(heat_maps[i][fsig_ind, :], axis=0)
        # fsig_all_med = np.median(heat_maps[i], axis=0)

        fsig_det = fsigs_concat_med

        ## umocnění feature signálu
        # fsig_det = (fsig_det ** 2)

        ## medianový filtr
        kernel_size = Config.kernel_size
        fsig_det_filt = medfilt(fsig_det, kernel_size=kernel_size)

        ## interpolace
        f = interpolate.interp1d(x_fsig, fsig_det_filt, kind='linear')
        fsig_det_filt_interp = f(x)

        ## uložení FS do listu
        fsigs.append(fsig_det_filt_interp)

    return fsigs

## ----------- Lokalizace paroxysmálních FiS -----------

def parox_afib_localization(fsigs, index_parox, test_df, qrs_loc_dict, loc_af_start_dict, loc_af_end_dict, ecg_data_dict, fs, signal_of_af_location_dict):

    det_start_all_list = list()
    det_end_all_list = list()
    det_start_all_true_list = list()
    det_end_all_true_list = list()

    for ii in range(len(fsigs)):

        ## ověření, že k feature signálu existuje v datech lokalizace - přeskočí se
        name_of_fsig = test_df.names[index_parox[ii]]
        if name_of_fsig not in signal_of_af_location_dict.keys():
            # print(f'Signal {name_of_fsig} has no ground truth localization.')
            continue

        ## ----- prahování -----
        fsig_to_thresh = fsigs[ii]
        threshold = np.std(fsig_to_thresh)
        # det_val = fsig_to_thresh[fsig_to_thresh > threshold]  # vrací hodnoty
        det_ind = np.where(fsig_to_thresh > threshold)[0]  # vrací indexy
        # det_bin = np.where(fsig_to_thresh > threshold, 1, 0)  # vrací binární vektor podle podmínky

        ## ----- získání start a end detekcí -----
        det_start_list = list()
        det_end_list = list()

        ## vytvoření detekcí
        for index in range(len(det_ind)):
            if index == 0:
                det_start_list.append(det_ind[0])
            elif index == (len(det_ind) - 1):
                det_end_list.append(det_ind[-1])
            else:
                if det_ind[index] - det_ind[index - 1] > 1:
                    det_start_list.append(det_ind[index])
                    det_end_list.append(det_ind[index - 1])

        ## omezení délky fibrilace ověřením počtu QRS komplexů - odstranění krátkých úseků
        det_start_qrs_list = list()
        det_end_qrs_list = list()
        qrs_loc = qrs_loc_dict[name_of_fsig]
        if (len(det_start_list) > 1) or (len(det_end_list) > 1):
            for q in range(len(det_start_list)):
                if len(qrs_loc[(qrs_loc >= det_start_list[q]) & (qrs_loc <= det_end_list[q])]) >= 5:
                    det_start_qrs_list.append(det_start_list[q])
                    det_end_qrs_list.append(det_end_list[q])
        else:
            det_start_qrs_list = det_start_list
            det_end_qrs_list = det_end_list

        det_start_edit_list = np.array(det_start_qrs_list)
        det_end_edit_list = np.array(det_end_qrs_list)

        ## omezení počtu fibrilací - uložení FiS s maximální hodnotou FS
        if (len(det_start_qrs_list) > 1) or (len(det_end_qrs_list) > 1):
            det_max = list()
            for m in range(len(det_start_qrs_list)):
                det_max.append(np.max(fsig_to_thresh[det_start_qrs_list[m]:det_end_qrs_list[m]]))
            arg_max = np.argmax(det_max)
            det_start_edit_list = np.array([det_start_qrs_list[arg_max]])
            det_end_edit_list = np.array([det_end_qrs_list[arg_max]])

        ## odstranění signálu, u kterého nebyla nalezena lokalizace
        if len(det_start_edit_list) == 0 and len(det_end_edit_list) == 0:
            # print(f'Localization of signal {name_of_fsig} has not been found!!!')
            continue
        ## uložení do listu
        det_start_all_list.extend(det_start_edit_list)
        det_end_all_list.extend(det_end_edit_list)

        ## ----- získání správných detekcích uložených v originálních datech
        det_start_true_list = [loc_af_start_dict[name_of_fsig]]
        det_end_true_list = [loc_af_end_dict[name_of_fsig]]

        ## uložení do listu
        det_start_all_true_list.extend(det_start_true_list)
        det_end_all_true_list.extend(det_end_true_list)

    return det_start_all_list, det_end_all_list, det_start_all_true_list, det_end_all_true_list

## Vyhodnocení lokalizace

def evaluation_of_localization(det_start_all_list, det_end_all_list, det_start_all_true_list, det_end_all_true_list, fs):

    ## porovnání detekcí - toleranční pásmo
    tolerance_zone_sec = 1  # dohromady jde o dvojnásobnou velikost tolerančního pásma např. 1 -> 2 s
    tolerance_zone_samples = fs * tolerance_zone_sec
    num_of_start_detections = len(det_start_all_list)
    num_of_end_detections = len(det_end_all_list)
    correct_start_detections = 0
    correct_end_detections = 0
    ## start
    for j in range(num_of_start_detections):
        if (det_start_all_true_list[j] - tolerance_zone_samples) <= det_start_all_list[j] <= (
                det_start_all_true_list[j] + tolerance_zone_samples):
            correct_start_detections += 1
    ## end
    for k in range(num_of_end_detections):
        if (det_end_all_true_list[k] - tolerance_zone_samples) <= det_end_all_list[k] <= (
                det_end_all_list[k] + tolerance_zone_samples):
            correct_end_detections += 1
    ## procento správných detekcí (signálů) v množině
    correct_start_detections_per = np.round((correct_start_detections / num_of_start_detections) * 100, 2)
    correct_end_detections_per = np.round((correct_end_detections / num_of_end_detections) * 100, 2)
    print(f'Start correct detections: {correct_start_detections_per} %')
    print(f'End correct detections: {correct_end_detections_per} %')


    ## porovnání detekcí - výpočet odchylky MAE
    mae_start_samples = np.abs(np.subtract(det_start_all_true_list, det_start_all_list))
    mae_start_sec_list = mae_start_samples / fs
    mae_start_sec = np.round(np.float32(mae_start_sec_list.mean(dtype=np.float32)), 2)
    sae_start_sec = np.round(np.float32(mae_start_sec_list.std(dtype=np.float32)), 2)
    print('Start MAE:', mae_start_sec)
    print('Start std AE:', sae_start_sec)
    mae_end_samples = np.abs(np.subtract(det_end_all_true_list, det_end_all_list))
    mae_end_sec_list = mae_end_samples / fs
    mae_end_sec = np.round(np.float32(mae_end_sec_list.mean(dtype=np.float32)), 2)
    sae_end_sec = np.round(np.float32(mae_end_sec_list.std(dtype=np.float32)), 2)
    print('End MAE:', mae_end_sec)
    print('End std AE:', sae_end_sec)
    # ## sloupcový graf - start
    # x = np.arange(len(mae_start_sec_list))
    # plt.figure()
    # plt.bar(x, mae_start_sec_list, width=1)
    # plt.xlabel('Pořadí signálu [-]')
    # plt.ylabel('Absolutní odchylka [s]')
    # plt.title('Identifikace začátků FiS')
    # plt.show()
    # ## sloupcový graf - end
    # x = np.arange(len(mae_end_sec_list))
    # plt.figure()
    # plt.bar(x, mae_end_sec_list, width=1)
    # plt.xlabel('Pořadí signálu [-]')
    # plt.ylabel('Absolutní odchylka [s]')
    # plt.title('Identifikace konců FiS')
    # plt.show()

    ## porovnání detekcí - výpočet odchylky
    err_start_samples = np.subtract(det_start_all_true_list, det_start_all_list)
    err_start_sec_list = err_start_samples / fs
    err_start_sec = np.round(np.float32(err_start_sec_list.mean(dtype=np.float32)), 2)
    serr_start_sec = np.round(np.float32(err_start_sec_list.std(dtype=np.float32)), 2)
    print('Start mean error:', err_start_sec)
    print('Start std error:', serr_start_sec)
    err_end_samples = np.subtract(det_end_all_true_list, det_end_all_list)
    err_end_sec_list = err_end_samples / fs
    err_end_sec = np.round(np.float32(err_end_sec_list.mean(dtype=np.float32)), 2)
    serr_end_sec = np.round(np.float32(err_end_sec_list.std(dtype=np.float32)), 2)
    print('End mean error:', err_end_sec)
    print('End std error:', serr_end_sec)
    # ## sloupcový graf - start
    # x = np.arange(len(err_start_sec_list))
    # plt.figure()
    # plt.bar(x, err_start_sec_list, width=1)
    # plt.xlabel('Pořadí signálu [-]')
    # plt.ylabel('Odchylka [s]')
    # plt.title('Identifikace začátků FiS')
    # plt.show()
    # ## sloupcový graf - end
    # x = np.arange(len(err_end_sec_list))
    # plt.figure()
    # plt.bar(x, err_end_sec_list, width=1)
    # plt.xlabel('Pořadí signálu [-]')
    # plt.ylabel('Odchylka [s]')
    # plt.title('Identifikace konců FiS')
    # plt.show()

    # return correct_start_detections_per, correct_end_detections_per, mse_start_sec, mse_end_sec, mae_start_sec, mae_end_sec, err_start_sec, err_end_sec