# import this file before anything using pytorch
import os
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
from typing import Iterable

import torch
import numpy as np

try:
    import hamiltorch
    __HAS_HAMILTORCH = True
except ImportError as e:
    __HAS_HAMILTORCH = False



def fix_seeds(seed: int = 42, warn_only=False):
    # for reproducibility on GPU with cudatoolkit >= 10.2
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

    # reproducibility
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    if not torch.are_deterministic_algorithms_enabled():
        torch.use_deterministic_algorithms(warn_only)

    if __HAS_HAMILTORCH:
        hamiltorch.set_random_seed(seed)
