"""
General utility functions for reproducibility, type normalization, tensor checks,
and model introspection used across training and data-processing pipelines.

Functions:
    set_seeds: Set random seeds for PyTorch CPU and CUDA operations.
    pair: Normalize input to tuple form of length 2.
    is_binary: Check if a tensor contains only binary values (0 or 1).
    model_size: Compute number of parameters and memory footprint of a model.
"""

import torch


def set_seeds(seed: int = 42):
    """
    Set random seeds for reproducibility in CPU and CUDA operations.
    """
    # Set the seed for general torch operations
    torch.manual_seed(seed)
    # Set the seed for CUDA torch operations (ones that happen on the GPU)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if using multi-GPU


def pair(t):
    """
    Ensure the input is a tuple of two identical elements if not already a tuple.
    """
    return t if isinstance(t, tuple) else (t, t)


def is_binary(x: torch.Tensor):
    """
    Returns True if the input tensor x contains only 0 and 1, False otherwise.
    """
    is_binary_vect = (x == 0) | (x == 1)
    return is_binary_vect.all()


def model_size(model):
    """
    Calculate the total number of parameters and memory footprint of a model.
    """
    n_params = sum(p.numel() for p in model.parameters())
    n_bytes = sum(p.element_size() * p.numel() for p in model.parameters())

    return n_params, n_bytes
