from abc import ABC, abstractmethod
from typing import Callable, Optional

import numpy as np


class RandomGenerator(ABC):
    @abstractmethod
    def _sample(self) -> np.ndarray:
        """The sampling function that used to generate random samples.
        :return: samples generated by defined sampling function
        :rtype: np.ndarray
        """
        pass

    def __call__(self) -> np.ndarray:
        """A function that calls sampling function and returns the generated sample.
        :return: samples generated by defined sampling function
        :rtype: np.ndarray
        """
        return self._sample()


class LambdaRandomGenerator(RandomGenerator):
    """
    LambdaRandomGenerator generates values by calling the supplied function.
    Note that instead of this you could also just use this supplied function itself.
    """

    def __init__(self, sampling_function: Callable[[], np.ndarray]):
        self.sampling_function = sampling_function

    def _sample(self) -> np.ndarray:
        return self.sampling_function()


class ReplayRandomGenerator(RandomGenerator):
    """
    ReplayRandomGenerator generates values by sequentially returning the values in the supplied
    array. It does not loop or reset automatically.
    """

    def __init__(self, values: np.ndarray):
        self._values = values
        self._idx = 0

    def _sample(self) -> np.ndarray:
        if self._idx == len(self._values):
            raise IndexError("ReplayRandomGenerator is out of values for index={}".format(self._idx))
        val = np.array(self._values[self._idx])
        self._idx += 1
        return val


class GaussianRandomGenerator(RandomGenerator):
    """
    GaussianRandomGenerator draws samples from a normal distribution with specified mean and standard deviation.
    Note that this gaussian can be multidimensional.
    """

    def __init__(self, mean: np.ndarray, std: np.ndarray, random_seed: Optional[int] = 0):
        self.mean = np.array(mean)
        self.std = np.array(std)
        self.rng = np.random.default_rng(random_seed)

    def _sample(self) -> np.ndarray:
        return self.rng.normal(self.mean, self.std)


class UniformRandomGenerator(RandomGenerator):
    """
    UniformRandomGenerator draws samples from a uniform distribution over the half-open interval [low, high).
    """

    def __init__(self, low: np.ndarray, high: np.ndarray):
        self.low = low
        self.high = high

    def _sample(self) -> np.ndarray:
        return np.random.uniform(self.low, self.high)
