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

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

import os
import numpy as np
import pandas as pd

import torch
from torch.utils.data import Dataset

# ###################################################
# ############### Class definition ##################
# ###################################################


# Torch dataset class for loading data and preparing (input,label) pairs for network model
class EKG_dataset(Dataset):

    # __init__ - constructor for initialization of dataset object
    def __init__(self, data_path, window_size, window_overlap, label_dict, shuffle=True):

        # Value initialization
        self.windows_size = window_size
        self.windows_shift = window_size - window_overlap
        self.patient_length = None
        self.shuffle = shuffle

        self.labels = []
        self.values = []
        self.locations = []

        # Patient loop
        for patient_id, patient_csv_path in enumerate(os.listdir(data_path)):

            if patient_csv_path[-3:] == "csv":
                patient_dataframe = pd.read_csv(os.path.join(data_path, patient_csv_path), sep=';')

                # Insert new empty lists for data of one patient
                self.labels.append([])
                self.values.append([])
                self.locations.append([])

                # Load patient labels
                for label_string in patient_dataframe['annotation'].values:
                    label_idx = label_dict[label_string]
                    # Insert new label value into patient list
                    self.labels[patient_id].append(label_idx)

                # Load patient data signals
                patient_data_orig = np.array(patient_dataframe['signal'].values, dtype=np.float32)
                patient_data_smooth = np.array(patient_dataframe['smooth'].values, dtype=np.float32)
                patient_data_deriv = np.array(patient_dataframe['derivation'].values, dtype=np.float32)

                # Set length of patient record
                if not self.patient_length:
                    self.patient_length = len(patient_data_orig)

                # Create samples from patient signals, where one sample is signle floating window 
                for i in range(0, self.patient_length, self.windows_shift):
                    orig_sample = patient_data_orig[i:i + self.windows_size]
                    smooth_sample = patient_data_smooth[i:i + self.windows_size]
                    deriv_sample = patient_data_deriv[i:i + self.windows_size]

                    sample = np.stack((orig_sample, smooth_sample, deriv_sample), axis=0)

                    # Ignore last window
                    if sample.shape[1] != self.windows_size:
                        break

                    self.values[patient_id].append(sample)

    # __len__(self) - return number of data set elements as (number of patients * number of samples per patient)
    def __len__(self):
        return len(self.values) * len(self.values[0])

    # __getitem__(self, idx) - return element on index "idx" from data set
    def __getitem__(self, idx):

        # At start of epoch shuffle patient data and labels 
        if self.shuffle:
            if idx == 0:
                self.shuffle_patients()

        # Compute from dataset idx, patient idx and window idx
        patient_idx = idx // len(self.values[0])
        window_idx = idx - patient_idx * len(self.values[0])

        # Select window labels from patient data
        start_idx = window_idx * self.windows_shift
        label_np = np.array(self.labels[patient_idx][start_idx:start_idx + self.windows_size])
        label_tensor = torch.tensor(label_np, dtype=torch.long)

        # Select window values and optionaly add spectrum to input from patient data
        value_np = np.array(self.values[patient_idx][window_idx], dtype=np.float32)

        value_tensor = torch.tensor(value_np, dtype=torch.float32)

        return value_tensor, label_tensor

    # shuffle patients - randomize order of patients data in dataset
    def shuffle_patients(self):

        # Create new ordering of patients
        new_order = np.random.permutation(range(len(self.values)))

        new_labels = []
        new_values = []

        # Shuffle patients
        for old_patient_id in new_order:
            new_labels.append(self.labels[old_patient_id])
            new_values.append(self.values[old_patient_id])

        # Rewrite old values
        self.labels = new_labels
        self.values = new_values
