#!/usr/bin/env python3
"""
Torch Determinism Shim (Basin Control)

Call `enable_determinism()` at process start to enforce:
- Deterministic CUDA operations
- Fixed seeds across all frameworks
- No randomness in sampling, dropout, or attention

Usage:
    from edge_app.torch_determinism import enable_determinism
    enable_determinism(seed=0)
"""
import os
import warnings


def enable_determinism(seed: int = 0, warn: bool = True):
    """
    Enable deterministic mode for PyTorch, TensorFlow, and NumPy.

    Args:
        seed: Fixed seed for all RNG sources
        warn: If True, warn about performance impact

    Raises:
        RuntimeError: If deterministic ops are not available
    """
    # Set environment variables BEFORE importing frameworks
    os.environ["PYTHONHASHSEED"] = str(seed)
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"  # Deterministic CUBLAS
    os.environ["TF_DETERMINISTIC_OPS"] = "1"           # TensorFlow determinism
    os.environ["TF_CUDNN_DETERMINISTIC"] = "1"         # TensorFlow cuDNN

    # PyTorch
    try:
        import torch
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.use_deterministic_algorithms(True, warn_only=False)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

        if warn:
            warnings.warn(
                "PyTorch deterministic mode enabled. Performance may be reduced (~10-30%). "
                "Some operations may raise errors if no deterministic implementation exists.",
                RuntimeWarning
            )
    except ImportError:
        pass  # PyTorch not installed
    except RuntimeError as e:
        raise RuntimeError(
            f"Failed to enable PyTorch determinism: {e}. "
            f"Check CUDA version and kernel support."
        ) from e

    # TensorFlow
    try:
        import tensorflow as tf
        tf.random.set_seed(seed)
        # Disable XLA auto-tuning (can introduce nondeterminism)
        os.environ["TF_XLA_FLAGS"] = "--tf_xla_auto_jit=0"

        if warn:
            warnings.warn(
                "TensorFlow deterministic mode enabled. Some ops may be slower.",
                RuntimeWarning
            )
    except ImportError:
        pass  # TensorFlow not installed

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

    # Random module
    import random
    random.seed(seed)


def verify_determinism():
    """
    Verify that determinism is enabled.

    Returns:
        dict: Status of determinism flags
    """
    status = {
        "pythonhashseed": os.environ.get("PYTHONHASHSEED"),
        "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG"),
        "tf_deterministic_ops": os.environ.get("TF_DETERMINISTIC_OPS"),
    }

    try:
        import torch
        status["torch_deterministic"] = torch.are_deterministic_algorithms_enabled()
        status["cudnn_deterministic"] = torch.backends.cudnn.deterministic
        status["cudnn_benchmark"] = torch.backends.cudnn.benchmark
    except ImportError:
        status["torch_deterministic"] = None

    try:
        import tensorflow as tf
        status["tf_seed"] = "set"  # No direct query API
    except ImportError:
        status["tf_seed"] = None

    return status


if __name__ == "__main__":
    # Test determinism
    enable_determinism(seed=42)
    status = verify_determinism()

    print("Determinism Status:")
    for key, val in status.items():
        print(f"  {key}: {val}")

    # Quick validation
    try:
        import torch
        x = torch.randn(10, device="cpu")
        y = torch.randn(10, device="cpu")
        print(f"\n✓ PyTorch determinism enabled")
        print(f"  Sample tensor: {x[0].item():.6f}")
    except ImportError:
        print("\n⚠ PyTorch not available")
