###### ------------------- augmentace dat pro DP -----------------

### importy
import numpy as np
import random
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from scipy import signal

from filereader import load_file

# ---------------------------- Dataset Transform -----------------------------

class Multiply():
    "vynásobení signálu danou hodnotou"

    ## vynásobení osy y
    def __init__(self, factor=None, p=0.3):
        self.factor = factor
        self.p = p

    ## vynásobení osy y hodnotou přibližně od 0.9 do 1.1
    def __call__(self, sample):
        inputs, targets = sample
        if torch.rand(1).numpy()[0] > self.p:
            if self.factor == None:
                # self.factor = torch.normal(1, 0.05, size=(1, 1))
                self.factor = round(random.uniform(0.9, 1.1), 2) ## 0.9-1.1 zaokrouhlené na 2 des. místa
            inputs *= self.factor
        return inputs, targets

class Add_Gauss():
    "přičtení Gaussovského šumu k signálu"

    def __init__(self, p=0.3):
        self.p = p

    def __call__(self, sample):
        inputs, targets = sample
        if torch.rand(1).numpy()[0] > self.p:
            # std = round(random.uniform(0, 0.02), 3) ## 0-0.02 zaokrouhlené na 3 des. místa
            noise = torch.normal(mean=0, std=0.1, size=inputs.size())
            inputs += noise
        return inputs, targets
