# - x - x - x - x - x - x - x - x - x - x - x - x - x - x - #
#                                                           #
#   This file was created by: Alberto Palomo Alonso         #
# Universidad de Alcalá - Escuela Politécnica Superior      #
#                                                           #
# - x - x - x - x - x - x - x - x - x - x - x - x - x - x - #
# Import statements:
import logging
import torch
import os
import numpy as np
import random
import time
# - x - x - x - x - x - x - x - x - x - x - x - x - x - x - #
CUBLAS_ALLOCATION = 4096
# - x - x - x - x - x - x - x - x - x - x - x - x - x - x - #


def get_seed(seed: int = None, logger: logging.Logger = None) -> int:
    """
    Sets the seed for generating random numbers to ensure reproducibility across numpy, random, and PyTorch operations.
    If no seed is provided, a new seed is generated based on the current time.

    This function also configures PyTorch to ensure deterministic behavior when running on a GPU, including the setting
    of environment variables to influence the behavior of CUDA's cuBLAS library.

    Args:
        seed (int, optional): The seed for the random number generators. If None, the seed will be generated based on
        the current system time.
        logger (logging.Logger): The logger that traces the logging information.

    Returns:
        int: The seed used to initialize the random number generators.

    Example:
        >>> experiment_seed = get_seed()
        Sets a random seed based on the current time and ensures that all subsequent random operations are reproducible.

        >>> experiment_seed = get_seed(42)
        >>> # experiment_seed == 42
        Uses 42 as the seed for all random number generators to ensure reproducibility.
    """
    # Set environment variable for deterministic behavior on CUDA >= 10.2
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = f":{CUBLAS_ALLOCATION}:8"

    # Create a new seed if not provided:
    seed = seed if seed is not None else int(time.time())

    # Set seed for numpy and random
    np.random.seed(seed)
    random.seed(seed)

    # Set seed and deterministic algorithms for torch
    torch.manual_seed(seed)
    torch.backends.cudnn.allow_tf32 = False
    torch.use_deterministic_algorithms(True, warn_only=True)

    # Ensure all operations are deterministic on GPU (if available)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    # Return the generated or bypassed seed:
    if logger is not None:
        logger.info(f"Initializer set up seed: {seed}")
    return seed
# - x - x - x - x - x - x - x - x - x - x - x - x - x - x - #
#                        END OF FILE                        #
# - x - x - x - x - x - x - x - x - x - x - x - x - x - x - #