"""Reproducibility utilities for training.

This module provides functions to set random seeds across all
libraries used in training (Python, NumPy, PyTorch) to ensure
reproducible results.

Example:
    >>> from essence_wars.training import set_seed
    >>> set_seed(42)  # Sets all seeds for reproducibility
    >>> set_seed(42, deterministic=True)  # Also enables CUDA determinism
"""

from __future__ import annotations

import os
import random


def set_seed(seed: int, deterministic: bool = False) -> None:
    """Set random seeds for reproducibility across all libraries.

    This function sets seeds for:
    - Python's random module
    - NumPy's random number generator
    - PyTorch's CPU and CUDA random number generators

    Args:
        seed: The random seed to use. Should be a non-negative integer.
        deterministic: If True, also enable CUDA deterministic algorithms.
            This may reduce performance but ensures reproducibility.
            Default is False for better training throughput.

    Example:
        >>> set_seed(42)
        >>> # Training will now be reproducible (modulo CUDA non-determinism)
        >>>
        >>> set_seed(42, deterministic=True)
        >>> # Full reproducibility, but potentially slower on GPU
    """
    # Python random
    random.seed(seed)

    # NumPy - imported lazily since it's always available
    import numpy as np

    np.random.seed(seed)

    # PyTorch - only if available (training extra)
    try:
        import torch

        torch.manual_seed(seed)

        if torch.cuda.is_available():
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)  # For multi-GPU

            if deterministic:
                # Enable deterministic algorithms
                torch.backends.cudnn.deterministic = True
                torch.backends.cudnn.benchmark = False

                # PyTorch 1.8+ deterministic flag
                if hasattr(torch, "use_deterministic_algorithms"):
                    torch.use_deterministic_algorithms(True)

                # Set environment variable for CUDA determinism
                os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

    except ImportError:
        # PyTorch not installed, skip
        pass


def get_seed_info() -> dict[str, int | str | None]:
    """Get information about current random state (for debugging).

    Returns:
        Dictionary with seed-related information from each library.
    """
    info: dict[str, int | str | None] = {}

    # Python random state (hash of state, not seed itself)
    info["python_random_state"] = hash(random.getstate()) % (2**31)

    # NumPy
    try:
        import numpy as np

        # Get first value from state for debugging
        state = np.random.get_state(legacy=False)
        if state["bit_generator"] == "MT19937":
            key = state["state"]["key"]
            info["numpy_mt_key_sample"] = int(key[0])
    except (ImportError, KeyError, TypeError):
        info["numpy"] = None

    # PyTorch
    try:
        import torch

        info["torch_seed"] = torch.initial_seed() % (2**31)
        if torch.cuda.is_available():
            info["cuda_available"] = True
            info["cudnn_deterministic"] = torch.backends.cudnn.deterministic
            info["cudnn_benchmark"] = torch.backends.cudnn.benchmark
        else:
            info["cuda_available"] = False
    except ImportError:
        info["torch"] = None

    return info
