###### ------------------- předzpracování dat a vytvoření nových datasetů DP -----------------

## slouží k načtení signálů a k vytvoření datasetu

### importy
import numpy as np
import torch
import matplotlib.pyplot as plt
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 Dataset, DataLoader
import random
import json
import glob
import os
from scipy import signal

from filereader import load_file
from config import Config
from transforms import Multiply, Add_Gauss

### ----------------------------------------------------------------------------------

def data_preprocessing(path, file_of_sample_names):

    ### načtení všech dat uvedených ve složce = datareader
    # path = Config.path_train

    ## vytvoření slovníku dat a labelů celého datasetu a rozdělení na 2 skupiny
    af_names = list()
    non_af_names = list()
    ecg_data_dict = {}
    labels_dict = {}
    labels_name_dict = {}
    qrs_loc_dict = {}
    loc_af_start_dict = {}
    loc_af_end_dict = {}
    signal_of_af_location_dict = {}

    length_of_af_list = list()

    ## rozdělení signálu - pomocné proměnné
    number_of_parox = 0
    number_of_persi = 0
    number_of_non = 0

    # file_of_sample_names = Config.file_of_sample_names
    file_set = open(os.path.join(path, file_of_sample_names), 'r').read().splitlines()
    for i, sample_name in enumerate(file_set):
        sample_path = os.path.join(path, sample_name)

        ## načtení signálu
        sig, sig_length, fs, label, global_labels, qrs_locations, fib_locations = load_file(sample_path)

        ## -------------- porovnání počtu signálů --------------
        if number_of_parox >= Config.max_number_of_parox and number_of_persi >= Config.max_number_of_persi and number_of_non >= Config.max_number_of_non:
            break

        if global_labels == 'paroxysmal atrial fibrillation' and number_of_parox >= Config.max_number_of_parox:
            continue
        elif global_labels == 'persistent atrial fibrillation' and number_of_persi >= Config.max_number_of_persi:
            continue
        elif global_labels == 'non atrial fibrillation' and number_of_non >= Config.max_number_of_non:
            continue

        ## ----------------- předzpracování dat -----------------

        ## filtrace driftu
        lowcut_drift = 0.5     ## mezní frekvence
        order_drift = 4         ## řád filtru
        sos_drift = signal.butter(order_drift, lowcut_drift, btype='highpass', fs=fs, output='sos') ## Butterworthův filtr, horní propust, mezní frekvence 0.5 Hz
        sig_filt_drift = signal.sosfiltfilt(sos_drift, sig, padtype='even', padlen=200)

        ## ----------------- rozdělení signálů - train = parox, persi a non ---------------------

        ## rozdělení signálu
        number_of_secs = Config.number_of_secs  # počet sekund pro získání počtu vzorků  např. 60 s
        number_of_samples = number_of_secs * fs  # počet vzorků, po kterých se signál rozdělí např. 12000 vzorků

        ## podmínka pro ověření, že délka signálu je větší než daný počet vzorků, aby nebyly signály kratší
        if sig_length < number_of_samples:
            continue

        factor = Config.factor
        ## vektor jednotlivých faktorů rozdělení od 0 do délky signálu s krokem number_of_samples
        vector_of_factors = np.arange(start=0, stop=sig_length, step=number_of_samples / factor, dtype=int)
        sig_start_ind = vector_of_factors[:-factor]
        sig_end_ind = vector_of_factors[factor:]

        ## vytvoření signálu pozic paroxysmálních fibrilací
        if global_labels == 'paroxysmal atrial fibrillation':

            af_start = [i for i in range(len(fib_locations)) if fib_locations[i] == '(AFIB' or fib_locations[i] == '(AFL']
            af_start_ind = qrs_locations[af_start]
            af_end = [i for i in range(len(fib_locations)) if fib_locations[i] == '(N']
            af_end_ind = qrs_locations[af_end]

            signal_of_af_location = np.zeros(sig_length)
            for s, e in zip(af_start_ind, af_end_ind):
                signal_of_af_location[s:e] = 1

            for s2, e2 in zip(af_start_ind, af_end_ind):
                length_of_af = e2 - s2  ## délka fibrilace
                ## uložení délky parox FiS do listu
                length_of_af_list.append(length_of_af)

                if length_of_af > number_of_samples:
                    continue

                for start, end in zip(sig_start_ind, sig_end_ind):
                    if s2 > start and e2 < end:
                        sig_div = sig_filt_drift[:, start:end]
                        signal_of_af_location_div = signal_of_af_location[start:end]
                        # break

                        if length_of_af == np.sum(signal_of_af_location_div):
                            ## vektor jednotlivých faktorů rozdělení od 0 do délky signálu s krokem number_of_samples
                            vector_of_factors = np.arange(start=end, stop=sig_length, step=number_of_samples / factor,
                                                          dtype=int)
                            sig_start_ind = vector_of_factors[:-factor]
                            sig_end_ind = vector_of_factors[factor:]

                            sample_name_new = sample_name + '_' + str(start) + '-' + str(end)

                            ## uložení pozic QRS komplexů
                            qrs_loc_dict[sample_name_new] = qrs_locations[(qrs_locations > start) & (qrs_locations < end)] - start

                            ## uložení pozice AFIB
                            loc_af_start_dict[sample_name_new] = s2 - start
                            loc_af_end_dict[sample_name_new] = e2 - start
                            signal_of_af_location_dict[sample_name_new] = signal_of_af_location_div

                            number_of_parox += 1

                            ## uložení ekg dat a labelů do slovníků
                            ecg_data_dict[sample_name_new] = sig_div
                            labels_dict[sample_name_new] = label
                            labels_name_dict[sample_name_new] = global_labels

                            ## rozdělení názvů dat podle labelů do listů
                            if labels_dict[sample_name_new] == 1:
                                af_names.append(sample_name_new)
                            elif labels_dict[sample_name_new] == 0:
                                non_af_names.append(sample_name_new)

                            break

        else:
            sig_div = sig_filt_drift[:, 0:number_of_samples]

            ## uložení ekg dat a labelů do slovníků
            ecg_data_dict[sample_name] = sig_div
            labels_dict[sample_name] = label
            labels_name_dict[sample_name] = global_labels

            ## rozdělení názvů dat podle labelů do listů
            if labels_dict[sample_name] == 1:
                af_names.append(sample_name)
                number_of_persi += 1

            elif labels_dict[sample_name] == 0:
                non_af_names.append(sample_name)
                number_of_non += 1

    af_names = np.array(af_names)
    non_af_names = np.array(non_af_names)

    return 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

def data_distribution(af_names, non_af_names, split_index_train=0.8, split_index_test=0.1):

    """
        Reference: https://stanford.edu/~shervine/blog/pytorch-how-to-generate-data-parallel
    """

    ## rozdělení dat na trénovací a validační -> náhodně pomocí indexování 80:20:
    temp_svt = np.arange(len(af_names))
    temp_non_svt = np.arange(len(non_af_names))
    random.seed(3)
    random.shuffle(temp_svt)
    random.shuffle(temp_non_svt)

    svt_split_index = np.ceil(split_index_train * len(af_names)).astype(int)
    non_svt_split_index = np.floor(split_index_train * len(non_af_names)).astype(int)
    svt_split_index_test = np.ceil(split_index_test * len(af_names)).astype(int)
    non_svt_split_index_test = np.floor(split_index_test * len(non_af_names)).astype(int)

    svt_train_index, svt_valid_index, svt_test_index = \
        temp_svt[:svt_split_index], \
        temp_svt[svt_split_index:svt_split_index + svt_split_index_test], \
        temp_svt[svt_split_index + svt_split_index_test:]
    non_svt_train_index, non_svt_valid_index, non_svt_test_index = \
        temp_non_svt[:non_svt_split_index], \
        temp_non_svt[non_svt_split_index:non_svt_split_index + non_svt_split_index_test], \
        temp_non_svt[non_svt_split_index + non_svt_split_index_test:]

    ## spojení svt a non_svt skupin názvů dat
    train_names, valid_names, test_names = list(), list(), list()

    train_names.extend(af_names[svt_train_index].tolist())
    valid_names.extend(af_names[svt_valid_index].tolist())
    test_names.extend(af_names[svt_test_index].tolist())

    train_names.extend(non_af_names[non_svt_train_index].tolist())
    valid_names.extend(non_af_names[non_svt_valid_index].tolist())
    test_names.extend(non_af_names[non_svt_test_index].tolist())

    train_names, valid_names, test_names = np.array(train_names), np.array(valid_names), np.array(test_names)

    temp_train_index, temp_valid_index, temp_test_index = np.arange(len(train_names)), np.arange(
        len(valid_names)), np.arange(len(test_names))
    random.seed(4)
    random.shuffle(temp_train_index), random.shuffle(temp_valid_index), random.shuffle(temp_test_index)

    partition = {"training": train_names[temp_train_index], "validation": valid_names[temp_valid_index],
                 "testing": test_names[temp_test_index]}

    return partition

### ----------------------------------------------------------------------------------