# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/98_utils.ipynb (unless otherwise specified).

__all__ = ['add_metrics_to_log', 'log_to_message', 'ProgressBar']

# Cell
#-------------------------------------------------------------------------------
# The ProgressBar, log_to_message, and add_metrics_to_log functions
# support Keras-style training progress in Jupyter notebooks.
# Code Reference:  https://github.com/henryre/pytorch-fitmodule
#-------------------------------------------------------------------------------
import numpy as np
import sys
def add_metrics_to_log(log, metrics, y_true, y_pred, prefix=''):
    for metric in metrics:
        q = metric(y_true, y_pred)
        log[prefix + metric.__name__] = q
    return log


def log_to_message(log, precision=4, omit=[]):
    fmt = "{0}: {1:." + str(precision) + "f}"
    fmt_fn = lambda k,v: "%s:%s" % (k,v) if k == 'iteration' else fmt.format(k,v)
    sep = "  "
    #return sep+sep.join(fmt.format(k, v) for k, v in log.items() if k not in omit)
    return sep+sep.join(fmt_fn(k, v) for k, v in log.items() if k not in omit)



class ProgressBar(object):
    """Progress bar"""

    def __init__(self, n, length=40):
        # Protect against division by zero
        self.n      = max(1, n)
        self.nf     = float(n)
        self.length = length
        # Precalculate the i values that should trigger a write operation
        self.ticks = set([round(i/100.0 * n) for i in range(101)])
        self.ticks.add(n-1)
        self.bar(0)

    def bar(self, i, message=""):
        """Assumes i ranges through [0, n-1]"""
        if i in self.ticks:
            b = int(np.ceil(((i+1) / self.nf) * self.length))
            sys.stdout.write("\r[{0}{1}] {2}%\t{3}".format(
                "="*b, " "*(self.length-b), int(100*((i+1) / self.nf)), message
            ))
            sys.stdout.flush()

    def close(self, message=""):
        # Move the bar to 100% before closing
        self.bar(self.n-1)
        sys.stdout.write("{0}\n\n".format(message))
        sys.stdout.flush()