import numpy as np

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

from spatz.transforms import Transform


class Downtime(Transform):
    def __init__(self, mu_duration: float, sigma_duration: float) -> None:
        super().__init__()

        self.__mu = mu_duration
        self.__sigma = sigma_duration
        self.__state = 1
        self.__until = abs(np.random.normal(mu_duration, sigma_duration))

    def get_state(self) -> Tuple[int, float]:
        return self.__state, self.__until

    def __call__(self, t: float, x: ArrayLike) -> Any:
        if t >= self.__until:
            self.__state = 1 - self.__state
            self.__until = t + abs(np.random.normal(self.__mu, self.__sigma))

        if self.__state == 1:
            return x
        
        return np.zeros_like(x)