"""Reproducibility utilities for deterministic abliteration runs.

Sets all random seeds (Python, NumPy, PyTorch CPU/CUDA) and optionally
enables PyTorch deterministic mode for bit-exact reproducibility.

Usage:
    from obliteratus.reproducibility import set_seed

    set_seed(42)  # sets all seeds, enables deterministic ops
    set_seed(42, deterministic=False)  # seeds only, faster
"""

from __future__ import annotations

import logging
import os
import random

logger = logging.getLogger(__name__)


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

    Args:
        seed: The seed value to use everywhere.
        deterministic: If True, also enable PyTorch deterministic algorithms
            and disable cuDNN benchmarking. This is slower but guarantees
            bit-exact results across runs on the same hardware.
    """
    random.seed(seed)

    try:
        import numpy as np
        np.random.seed(seed)
    except ImportError:
        pass

    try:
        import torch
        from obliteratus import device as dev
        torch.manual_seed(seed)
        dev.set_seed_all(seed)

        if deterministic:
            torch.use_deterministic_algorithms(True, warn_only=True)
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False
            os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    except ImportError:
        pass

    logger.debug("Seeds set to %d (deterministic=%s)", seed, deterministic)