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

from __future__ import annotations

__all__ = ["TorchRandomManager", "torch_seed"]

from contextlib import contextmanager
from typing import TYPE_CHECKING, Any

from coola.random.base import BaseRandomManager
from coola.utils.imports import check_torch, is_torch_available

if TYPE_CHECKING:
    from collections.abc import Generator

if is_torch_available():
    import torch
else:  # pragma: no cover
    from coola.utils.fallback.torch import torch


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

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

        ```
    """

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

    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) -> dict[str, Any]:
        return {
            "torch": torch.get_rng_state(),
            "torch.cuda": torch.cuda.get_rng_state_all(),
        }

    def manual_seed(self, seed: int) -> None:
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)

    def set_rng_state(self, state: dict[str, Any]) -> None:
        torch.set_rng_state(state["torch"])
        torch.cuda.set_rng_state_all(state["torch.cuda"])


@contextmanager
def torch_seed(seed: int) -> Generator[None, None, None]:
    r"""Implement a context manager to manage the PyTorch 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 torch
        >>> from coola.random import torch_seed
        >>> with torch_seed(42):
        ...     print(torch.randn(2, 4))
        ...
        tensor([[...]])
        >>> with torch_seed(42):
        ...     print(torch.randn(2, 4))
        ...
        tensor([[...]])

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