# utils.py
#
# some of the original utilities used by FFJORD
import os
import math
from numbers import Number
import logging
import torch


def makedirs(dirname):
    if not os.path.exists(dirname):
        os.makedirs(dirname)


def get_logger(logpath, filepath, package_files=[], displaying=True, saving=True, debug=False):
    logger = logging.getLogger()
    if debug:
        level = logging.DEBUG
    else:
        level = logging.INFO
    logger.setLevel(level)
    if saving:
        info_file_handler = logging.FileHandler(logpath, mode="a")
        info_file_handler.setLevel(level)
        logger.addHandler(info_file_handler)
    if displaying:
        console_handler = logging.StreamHandler()
        console_handler.setLevel(level)
        logger.addHandler(console_handler)
    logger.info(filepath)
    with open(filepath, "r") as f:
        logger.info(f.read())

    for f in package_files:
        logger.info(f)
        with open(f, "r") as package_f:
            logger.info(package_f.read())

    return logger


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


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
        self.sum = 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.sum += val
        self.val = val

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

#
# def inf_generator(iterable):
#     """Allows training with DataLoaders in a single infinite loop:
#         for i, (x, y) in enumerate(inf_generator(train_loader)):
#     """
#     iterator = iterable.__iter__()
#     while True:
#         try:
#             yield iterator.__next__()
#         except StopIteration:
#             iterator = iterable.__iter__()
#
#
# def save_checkpoint(state, save, epoch):
#     if not os.path.exists(save):
#         os.makedirs(save)
#     filename = os.path.join(save, 'checkpt-%04d.pth' % epoch)
#     torch.save(state, filename)
#
#
# def isnan(tensor):
#     return (tensor != tensor)
#
#
# def logsumexp(value, dim=None, keepdim=False):
#     """Numerically stable implementation of the operation
#     value.exp().sum(dim, keepdim).log()
#     """
#     if dim is not None:
#         m, _ = torch.max(value, dim=dim, keepdim=True)
#         value0 = value - m
#         if keepdim is False:
#             m = m.squeeze(dim)
#         return m + torch.log(torch.sum(torch.exp(value0), dim=dim, keepdim=keepdim))
#     else:
#         m = torch.max(value)
#         sum_exp = torch.sum(torch.exp(value - m))
#         if isinstance(sum_exp, Number):
#             return m + math.log(sum_exp)
#         else:
#             return m + torch.log(sum_exp)
#
#
#
