import numpy as np

from numpy.typing import ArrayLike
from typing import Any, Tuple

from spatz.transforms import Transform


class GaussianNoise(Transform):
    def __init__(self, mu: ArrayLike = None, sigma: ArrayLike = None) -> None:
        super().__init__()

        self.__mu = mu
        self.__sigma = sigma

    def __call__(self, _: float, x: ArrayLike) -> ArrayLike:
        if np.isscalar(x):
            noise = np.random.normal(0, 1)         

            x += self.__sigma * noise + self.__mu
        else:
            dim = len(x)

            if np.isscalar(self.__sigma):
                sigma = np.identity(dim) * self.__sigma
            else:
                sigma = self.__sigma
                
            if np.isscalar(self.__mu):
                mu = np.ones(dim)
            else:
                mu = self.__mu

            noise = np.random.normal(0, 1, np.shape(x))
            x += sigma @ noise + mu

        return x


class ProportionalGaussian(Transform):
    def __init__(self, mu, sigma) -> None:
        super().__init__()

        self.__mu = mu
        self.__sigma = sigma

    def __call__(self, _: float, x: ArrayLike) -> ArrayLike:
        noise = np.random.normal(0, 1)         

        x += (self.__sigma * x) * noise + (self.__mu * x)

        return x


class PinkNoise(Transform):
    def __init__(self) -> None:
        super().__init__()

    def __call__(self, t: float, x: ArrayLike) -> Any:
        pass