"""Utility functions to make your experiments reproducible

"""
import random
import warnings

import torch


def set_manual_seed(seed):
    """Set random seed for Python and PyTorch random generators.
    """
    random.seed(seed)
    torch.manual_seed(seed)

    print('Using manual seed: {seed}'.format(seed=seed))


def get_rng_state() -> dict:
    return {
        'torch_rng': torch.get_rng_state(),
        'torch_rng_cuda': torch.cuda.get_rng_state_all(),
        'python_rng': random.getstate(),
    }


def set_rng_state(rng_state: dict):
    try:
        torch_rng = rng_state['torch_rng']
        torch.set_rng_state(torch_rng)
        print('Set torch rng state')
    except ValueError as e:
        warnings.warn(e)

    try:
        torch_rng_cuda = rng_state['torch_rng_cuda']
        torch.cuda.set_rng_state(torch_rng_cuda)
        print('Set torch rng cuda state')
    except ValueError as e:
        warnings.warn(e)

    try:
        python_rng = rng_state['python_rng']
        random.setstate(python_rng)
        print('Set python rng state')
    except ValueError as e:
        warnings.warn(e)


def get_random_name() -> str:
    from . import namesgenerator as ng
    return ng.get_random_name()
