import torch
import random
import numpy as np


def set_seed(seed):
    """
    Seed all sources of randomness.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def seed_worker(worker_id):
    """
    Used to seed workers in a reproducible manner. Different seed per epoch. Initial seed is generated by
    the main process consuming one random number from the dataloader generator.

    trainloader = DataLoader(
        worker_init_fn=seed_worker,
        generator=loader_generator,
    )
    """

    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)
