###### ------------------- načítání dat pro jejich paralelizaci pro DP -----------------

### 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_reader(input, partition, data_dict, labels_dict):
    """
    input = 'training'/'validation'/'testing'
    """

    ### funkce augmentace dat
    composed = transforms.Compose([
        Multiply(factor=-1),  ## převrácení osy y
        Multiply(),  ## vynásobení osy y hodnotou od 0.9 do 1.1
        Add_Gauss()  ## přidání Gauss. šumu
    ])

    ### vytvoření datasetu
    if input == "training":
        output_dataset = Dataset(partition[input], data_dict, labels_dict, transform=composed)
    elif input == "validation":
        output_dataset = Dataset(partition[input], data_dict, labels_dict)
    elif input == "testing":
        output_dataset = Dataset(partition[input], data_dict, labels_dict)

    return output_dataset


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

    def __init__(self, list_IDs, data, labels, transform=None):
        "Inicializace"
        self.list_IDs = list_IDs
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        "Udává celkový počet vzorků"
        return len(self.list_IDs)

    def __getitem__(self, index):
        "Generuje jeden vzorek dat"
        ## výběr vzorku
        ID = self.list_IDs[index]

        ## načtení datového souboru a výběr labelu
        data = self.data[ID]  ## dict {ID: data}
        labels = self.labels[ID]  ## dict {ID: label}

        ## převedení na datový typ float32
        data = data.astype(np.float32)
        ## převedení na tenzor
        data = torch.from_numpy(data)

        ### one hot kodování labelu pro 2 třídy a softmax
        temp_labels = np.zeros(2, dtype=np.float32)
        temp_labels[labels] = 1
        labels = temp_labels
        # labels = labels.astype(np.float32)

        ## převedení na tenzor
        labels = torch.from_numpy(labels)

        ## augmentace dat - pouze training
        sample = data, labels
        if self.transform:
            sample = self.transform(sample)
        data, labels = sample

        return data, labels

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

### rozšíření signálů batche na stejnou délku
def collate_fn(data):
    """
    Reference: https://towardsdatascience.com/dataloader-for-sequential-data-using-pytorch-deep-learning-framework-part-2-ed3ad5f6ad82
                https://github.com/PetraNovotna/MIL_based_PVC_localization/blob/master/dataloader.py
    Input:
        data: list of tuple (training sequence, label)
    Output:
        padded_seqs - Padded Sequence, tensor of shape (batch_size, padded_length)
        labels - tensor of shape (batch_size)
        length - Original length of each sequence(without padding), tensor of shape(batch_size)
    """

    sequences, labels = zip(*data)

    lengths = [seq.shape[1] for seq in sequences]

    ## signály s kratší délkou se doplní nulami (padding) na délku nejdelšího signálu v batchi
    padded_seqs = np.zeros((len(sequences), sequences[0].shape[0], np.max(lengths))).astype(np.float32)
    for i, seq in enumerate(sequences):
        end = lengths[i]
        padded_seqs[i, :, :end] = seq

    ## převod na torch tensor
    padded_seqs = torch.from_numpy(padded_seqs)
    labels = np.stack(labels, axis=0)
    labels = torch.from_numpy(np.array(labels))
    lengths = np.array(lengths).astype(np.float32)
    lengths = torch.from_numpy(lengths)

    return padded_seqs, labels, lengths
### ----------------------------------------------------------------------------------