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) * self.__mu
            else:
                mu = self.__mu

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

        return x


class DriftingBias(Transform):
    def __init__(self, init: ArrayLike, covariance: ArrayLike, Tc: float) -> None:
        """First order Gauss-Markov (GM) model used to model drift.

        Args:
            init (ArrayLike): The initial bias.
            covariance (ArrayLike): Covariance matrix of the process.
            Tc (float): Correlation time of the process.
        """
        super().__init__()

        self.__t = 0
        self.__beta = 1 / Tc
        self.__covariance = covariance
        self.__Tc = Tc
        self.__x_old = np.copy(init)

    def __call__(self, t: float, x: ArrayLike) -> ArrayLike:
        dt = t - self.__t
        self.__t = t

        w = np.random.normal(np.zeros_like(x), self.__covariance*(1-np.exp(-2 * dt / self.__Tc)))
        drift = (1 - self.__beta * dt) * self.__x_old + w 
        
        self.__x_old = drift

        return x + drift


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