r"""Implement a random manager for NumPy."""

from __future__ import annotations

__all__ = ["NumpyRandomManager"]

from contextlib import contextmanager
from typing import TYPE_CHECKING, Any

from coola.random.base import BaseRandomManager
from coola.utils.imports import check_numpy, is_numpy_available

if TYPE_CHECKING:
    from collections.abc import Generator

if is_numpy_available():
    import numpy as np
else:  # pragma: no cover
    from coola.utils.fallback.numpy import numpy as np


class NumpyRandomManager(BaseRandomManager):  # noqa: PLW1641
    r"""Implement a random manager for the library ``numpy``.

    The seed must be between ``0`` and ``2**32 - 1``, so a modulo
    operator to convert an integer to an integer between ``0`` and
    ``2**32 - 1``.

    Example:
        ```pycon
        >>> from coola.random import NumpyRandomManager
        >>> manager = NumpyRandomManager()
        >>> manager.manual_seed(42)

        ```
    """

    def __init__(self) -> None:
        check_numpy()

    def __eq__(self, other: object) -> bool:
        return type(other) is type(self)

    def __repr__(self) -> str:
        return f"{self.__class__.__qualname__}()"

    def get_rng_state(self) -> tuple[Any, ...]:
        return np.random.get_state()

    def manual_seed(self, seed: int) -> None:
        np.random.seed(seed % 2**32)

    def set_rng_state(self, state: tuple[Any, ...]) -> None:
        np.random.set_state(state)


@contextmanager
def numpy_seed(seed: int) -> Generator[None, None, None]:
    r"""Implement a context manager to manage the NumPy random seed and
    random number generator (RNG) state.

    The context manager sets the specified random seed and
    restores the original RNG state afterward.

    Args:
        seed: The random number generator seed to use while using
            this context manager.

    Example:
        ```pycon
        >>> import numpy
        >>> from coola.random import numpy_seed
        >>> with numpy_seed(42):
        ...     print(numpy.random.randn(2, 4))
        ...
        [[...]]
        >>> with numpy_seed(42):
        ...     print(numpy.random.randn(2, 4))
        ...
        [[...]]

        ```
    """
    manager = NumpyRandomManager()
    state = manager.get_rng_state()
    try:
        manager.manual_seed(seed)
        yield
    finally:
        manager.set_rng_state(state)
