import random

import numpy as np
import torch
import time


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def timeit(method):
    def timed(*args, **kwargs):
        ts = time.time()
        result = method(*args, **kwargs)
        te = time.time()

        # first argument is `self`
        args[0].log_time[method.__name__] += [te - ts]
        return result

    return timed
