import numpy as np
import torch
import random

from .odenet_mnist.layers import MetaNODE
    
def fix_seeds(seed=502):
    np.random.seed(seed)
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.set_printoptions(precision=10)
    
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    
class RunningAverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self, momentum=0.99):
        self.momentum = momentum
        self.reset()

    def reset(self):
        self.val = None
        self.avg = 0

    def update(self, val):
        if self.val is None:
            self.avg = val
        else:
            self.avg = self.avg * self.momentum + val * (1 - self.momentum)
        self.val = val
        
        
def load_model(path):
    (_, state_dict), (_, model_args), (_, slover_id) = torch.load(path, map_location='cpu').items()

    is_odenet = model_args.network == 'odenet'
    
    if not hasattr(model_args, 'in_channels'):
        model_args.in_channels = 1
    
    model = MetaNODE(downsampling_method=model_args.downsampling_method,
                     is_odenet=is_odenet,
                     in_channels=model_args.in_channels)
    model.load_state_dict(state_dict) 
    
    return model, model_args
[((8, 4, 8, 24), 'numpy.random.seed', 'np.random.seed', ({(8, 19, 8, 23): 'seed'}, {}), '(seed)', True, 'import numpy as np\n'), ((9, 4, 9, 21), 'random.seed', 'random.seed', ({(9, 16, 9, 20): 'seed'}, {}), '(seed)', False, 'import random\n'), ((10, 4, 10, 27), 'torch.manual_seed', 'torch.manual_seed', ({(10, 22, 10, 26): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((11, 4, 11, 36), 'torch.cuda.manual_seed_all', 'torch.cuda.manual_seed_all', ({(11, 31, 11, 35): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((12, 4, 12, 40), 'torch.set_printoptions', 'torch.set_printoptions', (), '', False, 'import torch\n'), ((37, 55, 37, 91), 'torch.load', 'torch.load', (), '', False, 'import torch\n')]