from __future__ import annotations

from numpy.typing import ArrayLike
from numpy.random import Generator


class DeltaDistribution:
    def __init__(self, tau: float, rng: Generator | None = None):
        self._tau = tau
        self._rng = rng

    def __repr__(self):
        return f'DeltaDistribution (tau={self._tau})'

    def wait(self, size: int = 1) -> ArrayLike:
        return self._rng.exponential(self._tau, size=size)