# flake8: noqa
# This is taken (unedited) from keras:
#   https://github.com/keras-team/keras/blob/master/keras/utils/generic_utils.py#L285-L439
#   git sha: e6c3f77

import collections
import sys
import time

import numpy as np


class Progbar(object):
    """Displays a progress bar.

    Arguments:
        target (int): Total number of steps expected, None if unknown.
        width (float): Progress bar width on screen.
        verbose (int): Verbosity mode, 0 (silent), 1 (verbose), 2 (semi-verbose)
        stateful_metrics (list string): Iterable of string names of metrics that
            should *not* be averaged over time. Metrics in this list
            will be displayed as-is. All others will be averaged
            by the progbar before display.
        interval (float): Minimum visual progress update interval (in seconds).
        alpha (float): The coefficient for exponentially weighted moving averages.
    """

    def __init__(self, target, width=30, verbose=1, interval=0.05,
                 stateful_metrics=None, alpha=0.05):
        self.target = target
        self.width = width
        self.verbose = verbose
        self.interval = interval

        if stateful_metrics:
            self.stateful_metrics = set(stateful_metrics)
        else:
            self.stateful_metrics = set()

        self.alpha = alpha

        self._dynamic_display = ((hasattr(sys.stdout, 'isatty') and
                                  sys.stdout.isatty()) or
                                 'ipykernel' in sys.modules)
        self._total_width = 0
        self._seen_so_far = 0
        self._avg_values = {}
        self._start = time.time()
        self._last_update = 0

    def update(self, current, values=None):
        """Updates the progress bar.

        Arguments:
            current (int): Index of current step.
            values (list tuples): List of tuples like ``(name, value_for_last_step)``.

        Note:
            If ``name`` is in ``stateful_metrics``, ``value_for_last_step`` will be displayed as-is.
            Else, an average of the metric over time will be displayed.
        """
        values = values or []
        for k, v in values:
            if k not in self.stateful_metrics:
                if k not in self._avg_values:
                    self._avg_values[k] = v
                else:
                    self._avg_values[k] = (self.alpha * v) + ((1 - self.alpha) * self._avg_values[k])
        self._seen_so_far = current

        now = time.time()
        info = ' - %.0fs' % (now - self._start)
        if self.verbose == 1:
            if (now - self._last_update < self.interval and
                    self.target is not None and current < self.target):
                return

            prev_total_width = self._total_width
            if self._dynamic_display:
                sys.stdout.write('\b' * prev_total_width)
                sys.stdout.write('\r')
            else:
                sys.stdout.write('\n')

            if self.target is not None:
                numdigits = int(np.floor(np.log10(self.target))) + 1
                barstr = '%%%dd/%d [' % (numdigits, self.target)
                bar = barstr % current
                prog = float(current) / self.target
                prog_width = int(self.width * prog)
                if prog_width > 0:
                    bar += ('=' * (prog_width - 1))
                    if current < self.target:
                        bar += '>'
                    else:
                        bar += '='
                bar += ('.' * (self.width - prog_width))
                bar += ']'
            else:
                bar = '%7d/Unknown' % current

            self._total_width = len(bar)
            sys.stdout.write(bar)

            if current:
                time_per_unit = (now - self._start) / current
            else:
                time_per_unit = 0
            if self.target is not None and current < self.target:
                eta = time_per_unit * (self.target - current)
                if eta > 3600:
                    eta_format = '%d:%02d:%02d' % (eta // 3600, (eta % 3600) // 60, eta % 60)
                elif eta > 60:
                    eta_format = '%d:%02d' % (eta // 60, eta % 60)
                else:
                    eta_format = '%ds' % eta

                info = ' - ETA: %s' % eta_format
            else:
                if time_per_unit >= 1:
                    info += ' %.0fs/step' % time_per_unit
                elif time_per_unit >= 1e-3:
                    info += ' %.0fms/step' % (time_per_unit * 1e3)
                else:
                    info += ' %.0fus/step' % (time_per_unit * 1e6)

            for name, avg_value in self._avg_values.items():
                info += ' - %s:' % name
                if abs(avg_value) > 1e-3:
                    info += ' %.4f' % avg_value
                else:
                    info += ' %.4e' % avg_value

            self._total_width += len(info)
            if prev_total_width > self._total_width:
                info += (' ' * (prev_total_width - self._total_width))

            if self.target is not None and current >= self.target:
                info += '\n'

            sys.stdout.write(info)
            sys.stdout.flush()

        elif self.verbose == 2:
            if self.target is None or current >= self.target:
                for name, avg_value in self._avg_values.items():
                    info += ' - %s:' % name
                    if abs(avg_value) > 1e-3:
                        info += ' %.4f' % avg_value
                    else:
                        info += ' %.4e' % avg_value
                info += '\n'

                sys.stdout.write(info)
                sys.stdout.flush()

        self._last_update = now

    def add(self, n, values=None):
        self.update(self._seen_so_far + n, values)
