"""Reproducibility utilities for ML experiments.

This module provides helpers for setting random seeds across all
common ML frameworks to ensure reproducible experiments.

Author: Generated by AgentBible
"""

import os
import random
import sys
from dataclasses import dataclass
from typing import Any, Optional


@dataclass
class SeedState:
    """Container for random state across all frameworks."""

    seed: int
    numpy_state: Optional[Any] = None
    python_state: Optional[Any] = None
    torch_state: Optional[Any] = None
    torch_cuda_state: Optional[Any] = None


def set_all_seeds(seed: int = 42, deterministic: bool = True) -> None:
    """Set random seeds for all common ML frameworks.

    Sets seeds for: Python random, NumPy, PyTorch, TensorFlow, and JAX.
    Also sets environment variables for additional determinism.

    Args:
        seed: Random seed value.
        deterministic: If True, enable deterministic algorithms (may be slower).

    Example:
        >>> set_all_seeds(42)  # Call at start of experiment
    """
    # Python random
    random.seed(seed)

    # NumPy
    import numpy as np

    np.random.seed(seed)

    # Environment variables for reproducibility
    os.environ["PYTHONHASHSEED"] = str(seed)

    # PyTorch (if available)
    try:
        import torch

        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)
        if deterministic:
            torch.backends.cudnn.deterministic = True  # type: ignore[attr-defined]
            torch.backends.cudnn.benchmark = False  # type: ignore[attr-defined]
    except ImportError:
        pass

    # TensorFlow (if available)
    try:
        import tensorflow as tf

        tf.random.set_seed(seed)
        if deterministic:
            os.environ["TF_DETERMINISTIC_OPS"] = "1"
    except ImportError:
        pass

    # JAX (if available) - JAX uses explicit key-based randomness
    # but we can set the numpy seed for any numpy operations
    try:
        import jax

        # JAX doesn't have global seed, but document the intended key
        _ = jax.random.PRNGKey(seed)  # User should use this key
    except ImportError:
        pass


def get_seed_state(seed: int = 42) -> SeedState:
    """Capture current random state for later restoration.

    Args:
        seed: The seed that was used.

    Returns:
        SeedState object containing all random states.
    """
    import numpy as np

    state = SeedState(
        seed=seed,
        numpy_state=np.random.get_state(),
        python_state=random.getstate(),
    )

    try:
        import torch

        state.torch_state = torch.get_rng_state()
        if torch.cuda.is_available():
            state.torch_cuda_state = torch.cuda.get_rng_state_all()
    except ImportError:
        pass

    return state


def restore_seed_state(state: SeedState) -> None:
    """Restore random state from a SeedState object.

    Args:
        state: Previously captured SeedState.
    """
    import numpy as np

    if state.numpy_state is not None:
        np.random.set_state(state.numpy_state)
    if state.python_state is not None:
        random.setstate(state.python_state)

    try:
        import torch

        if state.torch_state is not None:
            torch.set_rng_state(state.torch_state)
        if state.torch_cuda_state is not None and torch.cuda.is_available():
            torch.cuda.set_rng_state_all(state.torch_cuda_state)
    except ImportError:
        pass


def get_experiment_config(
    seed: int = 42,
    extra_params: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    """Generate a reproducibility config for experiment tracking.

    Args:
        seed: Random seed used.
        extra_params: Additional parameters to include.

    Returns:
        Dictionary suitable for logging to wandb/mlflow.
    """
    import numpy as np

    config: dict[str, Any] = {
        "seed": seed,
        "numpy_version": np.__version__,
        "python_version": sys.version,
    }

    try:
        import torch

        config["torch_version"] = torch.__version__
        config["cuda_available"] = torch.cuda.is_available()
        if torch.cuda.is_available():
            config["cuda_version"] = torch.version.cuda
    except ImportError:
        pass

    try:
        import sklearn

        config["sklearn_version"] = sklearn.__version__
    except ImportError:
        pass

    if extra_params:
        config.update(extra_params)

    return config
