# -*- encoding: utf-8 -*-
"""
Text progress bar for long-running jobs.

:copyright: (c) 2016 H2O.ai
:license:   Apache License Version 2.0 (see LICENSE for details)
"""
from __future__ import absolute_import, division, print_function, unicode_literals

import math
import os
import sys
import time
from types import FunctionType, GeneratorType, MethodType

import colorama
from h2o.utils.compatibility import *  # NOQA
from h2o.utils.shared_utils import clamp
from h2o.utils.typechecks import assert_is_type, is_type, numeric

class ProgressBar(object):
    """
    Component that controls execution of a remote job, and draws a progress bar as it goes along.

    Usage example::

        def progress_generator():
            for i in range(10):
                yield i + 1, 0.5

        from progressbar import ProgressBar, PBWBar, PBWPercentage
        pb = ProgressBar(maxval=10, widgets_tty=["Job progress:", PBWBar(), PBWPercentage()])
        pb.execute(progress_generator())

    One of the important features of this class is that it takes an active role in controlling the progress of the
    underlying process. In particular, the :meth:`execute` method takes a ``progress_fn`` parameter, which is a
    generator that the :meth:`execute` method will periodically call in order to find what the progress level is.
    The generator may return either a single value (the progress amount), or a tuple
    ``(progress_amount, poll_interval)``, where the second element specifies how long (in seconds) ProgressBar
    should wait before calling the generator again.

    Once the :meth:`execute` method is called, it will not return until the progress reaches 100%. However the progress
    function can raise ``StopIteration(message)`` in order to cancel the current progress (and if possible display the
    provided message, which should be short).

    Another feature of this class is that it models the progress as a smooth process, so that progress values reported
    by the executor are used merely as signals, and the class makes up its own mind regarding what the actual progress
    at each moment of time is. More specifically, this works as follows:

    Suppose x(t) is the progress level as a function of time t, and v(t) is its speed (i.e. v(t) = dx/dt). Let T be
    the expected time when progress reaches 100%. Suppose current progress level is x₀ and speed v₀ at time t₀. Ideally
    we'd want to move to x(T) = 1 at a constant speed, which would be (1-x₀)/(T-t₀). However unless current v₀ is
    exactly that, such strategy is not feasible (we do not allow sudden jumps in speed). Instead, we let the
    progress move with the speed v(t) = vₑ + (v₀ - vₑ)e⁻ᵝ⁽ᵗ⁻ᵗ⁰⁾, where vₑ is the target speed that has to be
    determined based on the requirement to reach 1 at time T, and β is the rate of convergence towards the target
    speed vₑ.

    After some simple calculus, we obtain
        x(t) = x₀ + vₑ·(t - t₀) + (v₀ - vₑ)/β·(1 - e⁻ᵝ⁽ᵗ⁻ᵗ⁰⁾),
        v(t) = vₑ + (v₀ - vₑ)e⁻ᵝ⁽ᵗ⁻ᵗ⁰⁾,  and
        vₑ = v₀ + β(1 - x₀ - v₀(T - t₀))/(β(T - t₀) - 1 + e⁻ᵝ⁽ᵀ⁻ᵗ⁰⁾)

    In order to ensure smoothness of the model, we need to be careful as the progress approaches 100% and t
    approaches T. First, under no circumstance we are allowed to finish the progress if the underlying process did
    not report 100% status. At the same time we don't want to jump suddenly from, say, 90% to 100% at the final step.
    Thus, we allow the model to progress for a very short amount of time after the process has completed.
    """

    # Minimum and maximum frequency for progress checks (i.e. do not query progress faster than every
    # MIN_PROGRESS_CHECK_INTERVAL seconds, or slower than every MAX_PROGRESS_CHECK_INTERVAL seconds).
    MIN_PROGRESS_CHECK_INTERVAL = 0.2
    MAX_PROGRESS_CHECK_INTERVAL = 5.0

    # Model parameter: rate of speed adjustment (in 1/s)
    BETA = 0.3

    # This parameter determines calculation of local speed progress.
    GAMMA = 0.6

    # How long are we allowed to wait after the progress is finished before returning.
    FINISH_DELAY = 0.3

    def __init__(self, title=None, widgets=None, maxval=1.0, file_mode=None, hidden=False):
        """
        Initialize the progress bar.

        :param title: name of the progress bar, to be used as the first default widget.
        :param maxval: when the progress reaches this value, the job is considered finished.
        :param widgets: list of widgets to render. Each widget is either a constant string, or an
            instance of a ProgressBarWidget subclass. Defaults are: [title+":", PBWBar(), PBWPercentage()].
        :param file_mode: if True, force file-mode output (i.e. progress is printed without ability to
            overwrite previous output).
        """
        assert_is_type(title, None, str)
        assert_is_type(maxval, numeric)
        assert_is_type(widgets, None, [str, ProgressBarWidget])
        assert_is_type(file_mode, None, bool)

        # Fix for PUBDEV-5048. H2O depends on isatty attribute, but some Python Notebooks override stdout and doesn't
        # specify it. The same holds for the encoding attribute bellow
        if not hasattr(sys.stdout, "isatty"):
            sys.stdout.isatty = lambda: False

        if not hasattr(sys.stdout, "encoding"):
            sys.stdout.encoding = sys.getdefaultencoding()

        if title is None: title = "Progress"
        if file_mode is None: file_mode = not sys.stdout.isatty()

        self._maxval = maxval
        self._file_mode = file_mode
        if hidden:
            self._widget = _HiddenWidget()
            self._file_mode = True
        else:
            self._widget = _ProgressBarCompoundWidget(widgets, title=title, file_mode=file_mode)

        # Variables needed for progress model (see docs).
        self._t0 = None  # Time when the model's parameters were computed
        self._x0 = None  # Progress level at t0, a real number from 0 to 1
        self._v0 = None  # Progress speed at t0, in 1/s
        self._ve = None  # Target progress speed

        # List of tuples (timestamp, progress_value) reported by the external source. The timestamps are taken from
        # the moment the data was received, not when it was requested. The progress values are raw (i.e. they range
        # from 0 to self._maxval).
        self._progress_data = []
        # Timestamp when should the progress be queried next.
        self._next_poll_time = None

    def execute(self, progress_fn, print_verbose_info=None):
        """
        Start the progress bar, and return only when the progress reaches 100%.

        :param progress_fn: the executor function (or a generator). This function should take no arguments
            and return either a single number -- the current progress level, or a tuple (progress level, delay),
            where delay is the time interval for when the progress should be checked again. This function may at
            any point raise the ``StopIteration(message)`` exception, which will interrupt the progress bar,
            display the ``message`` in red font, and then re-raise the exception.
        :raises StopIteration: if the job is interrupted. The reason for interruption is provided in the exception's
            message. The message will say "cancelled" if the job was interrupted by the user by pressing Ctrl+C.
        """
        assert_is_type(progress_fn, FunctionType, GeneratorType, MethodType)
        if isinstance(progress_fn, GeneratorType):
            # Convert generator to a regular function
            progress_fn = (lambda g: lambda: next(g))(progress_fn)

        # Initialize the execution context
        self._next_poll_time = 0
        self._t0 = time.time()
        self._x0 = 0
        self._v0 = 0.01  # corresponds to 100s completion time
        self._ve = 0.01

        progress = 0
        status = None  # Status message in case the job gets interrupted.
        try:
            while True:
                # We attempt to synchronize all helper functions, ensuring that each of them has the same idea
                # for what the current time moment is. Otherwise we could have some corner cases when one method
                # says that something must happen right now, while the other already sees that moment in the past.
                now = time.time()

                # Query the progress level, but only if it's time already
                if self._next_poll_time <= now:
                    res = progress_fn()  # may raise StopIteration
                    assert_is_type(res, (numeric, numeric), numeric)
                    if not isinstance(res, tuple):
                        res = (res, -1)
                    # Progress querying could have taken some time, so update the current time moment
                    now = time.time()
                    self._store_model_progress(res, now)
                    self._recalculate_model_parameters(now)

                # Render the widget regardless of whether it's too early or not
                progress = min(self._compute_progress_at_time(now)[0], 1)
                if progress == 1 and self._get_real_progress() >= 1:
                    # Do not exit until both the model and the actual progress reach 100% mark.
                    break
                result = self._widget.render(progress)
                assert_is_type(result, RenderResult)
                time0 = result.next_time
                time1 = self._get_time_at_progress(result.next_progress)
                next_render_time = min(time0, time1)
                self._draw(result.rendered)

                # Wait until the next rendering/querying cycle
                wait_time = min(next_render_time, self._next_poll_time) - now
                if wait_time > 0:
                    time.sleep(wait_time)
                    if print_verbose_info is not None:
                        print_verbose_info(progress)
        except KeyboardInterrupt:
            # If the user presses Ctrl+C, we interrupt the progress bar.
            status = "cancelled"
        except StopIteration as e:
            # If the generator raises StopIteration before reaching 100%, then the progress display will
            # reamin incomplete.
            status = str(e)

        # Do one final rendering before we exit
        result = self._widget.render(progress=progress, status=status)
        self._draw(result.rendered, final=True)

        if status == "cancelled":
            # Re-raise the exception, to inform the upstream caller that something unexpected happened.
            raise StopIteration(status)

    #-------------------------------------------------------------------------------------------------------------------
    #  Private
    #-------------------------------------------------------------------------------------------------------------------

    def _get_real_progress(self):
        return self._progress_data[-1][1] / self._maxval

    def _store_model_progress(self, res, now):
        """
        Save the current model progress into ``self._progress_data``, and update ``self._next_poll_time``.

        :param res: tuple (progress level, poll delay).
        :param now: current timestamp.
        """
        raw_progress, delay = res
        raw_progress = clamp(raw_progress, 0, self._maxval)
        self._progress_data.append((now, raw_progress))

        if delay < 0:
            # calculation of ``_guess_next_poll_interval()`` should be done only *after* we pushed the fresh data to
            # ``self._progress_data``.
            delay = self._guess_next_poll_interval()
        self._next_poll_time = now + clamp(delay, self.MIN_PROGRESS_CHECK_INTERVAL, self.MAX_PROGRESS_CHECK_INTERVAL)

    def _recalculate_model_parameters(self, now):
        """Compute t0, x0, v0, ve."""
        time_until_end = self._estimate_progress_completion_time(now) - now
        assert time_until_end >= 0, "Estimated progress completion cannot be in the past."
        x_real = self._get_real_progress()
        if x_real == 1:
            t0, x0, v0, ve = now, 1, 0, 0
        else:
            x0, v0 = self._compute_progress_at_time(now)
            t0 = now
            if x0 >= 1:
                # On rare occasion, the model's progress may have reached 100% by ``now``. This can happen if
                # (1) the progress is close to 100% initially and has high speed, (2) on the previous call we
                # estimated that the process completion time will be right after the next poll time, and (3)
                # the polling itself took so much time that the process effectively "overshoot".
                # If this happens, then we adjust x0, v0 to the previous valid data checkpoint.
                t0, x0, v0 = self._t0, self._x0, self._v0
                time_until_end += now - t0
            z = self.BETA * time_until_end
            max_speed = (1 - x_real**2) / self.FINISH_DELAY
            ve = v0 + (self.BETA * (1 - x0) - v0 * z) / (z - 1 + math.exp(-z))
            if ve < 0:
                # Current speed is too high -- reduce v0 (violate non-smoothness of speed)
                v0 = self.BETA * (1 - x0) / (1 - math.exp(-z))
                ve = 0
            if ve > max_speed:
                # Current speed is too low: finish later, but do not allow ``ve`` to be higher than ``max_speed``
                ve = max_speed
        self._t0, self._x0, self._v0, self._ve = t0, x0, v0, ve

    def _estimate_progress_completion_time(self, now):
        """
        Estimate the moment when the underlying process is expected to reach completion.

        This function should only return future times. Also this function is not allowed to return time moments less
        than self._next_poll_time if the actual progress is below 100% (this is because we won't know that the
        process have finished until we poll the external progress function).
        """
        assert self._next_poll_time >= now
        tlast, wlast = self._progress_data[-1]
        # If reached 100%, make sure that we finish as soon as possible, but maybe not immediately
        if wlast == self._maxval:
            current_completion_time = (1 - self._x0) / self._v0 + self._t0
            return clamp(current_completion_time, now, now + self.FINISH_DELAY)

        # Calculate the approximate speed of the raw progress based on recent data
        tacc, wacc = 0, 0
        factor = self.GAMMA
        for t, x in self._progress_data[-2::-1]:
            tacc += factor * (tlast - t)
            wacc += factor * (wlast - x)
            factor *= self.GAMMA
            if factor < 1e-2: break

        # If there was no progress at all, then just assume it's 5 minutes from now
        if wacc == 0: return now + 300

        # Estimate the completion time assuming linear progress
        t_estimate = tlast + tacc * (self._maxval - wlast) / wacc

        # Adjust the estimate if it looks like it may happen too soon
        if t_estimate <= self._next_poll_time:
            t_estimate = self._next_poll_time + self.FINISH_DELAY

        return t_estimate

    def _guess_next_poll_interval(self):
        """
        Determine when to query the progress status next.

        This function is used if the external progress function did not return time interval for when it should be
        queried next.
        """
        time_elapsed = self._progress_data[-1][0] - self._progress_data[0][0]
        real_progress = self._get_real_progress()
        return min(0.2 * time_elapsed, 0.5 + (1 - real_progress)**0.5)

    def _compute_progress_at_time(self, t):
        """
        Calculate the modelled progress state for the given time moment.

        :returns: tuple (x, v) of the progress level and progress speed.
        """
        t0, x0, v0, ve = self._t0, self._x0, self._v0, self._ve
        z = (v0 - ve) * math.exp(-self.BETA * (t - t0))
        vt = ve + z
        xt = clamp(x0 + ve * (t - t0) + (v0 - ve - z) / self.BETA, 0, 1)
        return xt, vt

    def _get_time_at_progress(self, x_target):
        """
        Return the projected time when progress level `x_target` will be reached.

        Since the underlying progress model is nonlinear, we need to do use Newton method to find a numerical solution
        to the equation x(t) = x_target.
        """
        t, x, v = self._t0, self._x0, self._v0
        # The convergence should be achieved in just few iterations, however in unlikely situation that it doesn't
        # we don't want to loop forever...
        for _ in range(20):
            if v == 0: return 1e20
            # make time prediction assuming the progress will continue at a linear speed ``v``
            t += (x_target - x) / v
            # calculate the actual progress at that time
            x, v = self._compute_progress_at_time(t)
            # iterate until convergence
            if abs(x - x_target) < 1e-3: return t
        return time.time() + 100

    def _draw(self, txt, final=False):
        """Print the rendered string to the stdout."""
        if not self._file_mode:
            # If the user presses Ctrl+C this ensures we still start writing from the beginning of the line
            sys.stdout.write("\r")
        sys.stdout.write(txt)
        if final and not isinstance(self._widget, _HiddenWidget):
            sys.stdout.write("\n")
        else:
            if not self._file_mode:
                sys.stdout.write("\r")
            sys.stdout.flush()

    def __repr__(self):
        """Progressbar internal state (for debug purposes)."""
        t0 = self._progress_data[0][0]
        data = ",".join("(%.1f,%.3f)" % (t - t0, w / self._maxval) for t, w in self._progress_data)
        return "<Progressbar x0=%.3f, v0=%.3f, xraw=%.3f; data:[%s]>" % \
            (self._x0, self._v0, self._get_real_progress(), data)

#-----------------------------------------------------------------------------------------------------------------------
# Widgets (abstract)
#-----------------------------------------------------------------------------------------------------------------------

class RenderResult(object):
    """
    Helper class which serves as a return value from `ProgressBarWidget.render` method.

    It consists of the rendered string itself, and the schedule for when the widget has to be rendered again. The
    widget can request to be rerendered at a specific moment in time, or when the process reaches a certain progress
    level (or whichever comes sooner).
    """

    def __init__(self, rendered="", length=None, next_progress=1.0, next_time=1e20):
        """
        Initialize a new RenderResult instance.

        :param rendered: string representing the rendered widget.
        :param length: the character length of the rendered result (this could be different from ``len(rendered)`` if
            the result contains ANSI escape sequences).
        :param next_progress: request the widget to be rendered again when process' progress reaches this level.
        :param next_time: request the widget to be rendered again at this moment in time.
        """
        assert_is_type(rendered, str)
        assert_is_type(length, int, None)
        assert_is_type(next_progress, float)
        assert_is_type(next_time, int, float)
        self.rendered = rendered
        self.length = length if length is not None else len(rendered)
        self.next_progress = next_progress
        self.next_time = next_time

class ProgressBarWidget(object):
    """
    Base class for all progress bar widgets.

    The primary interface of this class is the ``render(progress)`` method, which requests the widget to render
    itself for the given progress level. The method returns a ``RenderResult`` object, consisting of: the rendered
    string itself, the next time moment when rendering should occur, or the next progress level at which the widget
    should be rerendered.
    """

    def __init__(self):
        """Initialize a progress bar widget."""
        self._file_mode = False

    def render(self, progress, width=None, status=None):
        """
        Render the widget.

        :param float progress: current job progress, a number between 0 and 1.
        :param int width: target character widths (for flexible widgets).
        :param str status: if the job has finished unexpectedly, this will be a message explaining the reason.
        :returns: a ``RenderResult`` object.
        """
        raise NotImplementedError()

    def set_encoding(self, encoding):
        """
        Inform the widget about the character encoding of the underlying output stream.

        Some widgets may render differently depending on whether the output is unicode-aware or not. This function
        will be called on the widget before its first render, to tell the widget what the encoding will be. It is up
        to the widget whether to do anything with this information or simply ignore it (default).

        :param encoding: the encoding string, eg "utf-8" or "cp437".
        """
        assert_is_type(encoding, str)

    def set_mode(self, mode):
        """
        Inform the widget that it will be rendered in either tty or file mode.

        This is only useful for widgets that support dual rendering mode.

        :param mode: either "tty" or "file".
        """
        assert_is_type(mode, "tty", "file")
        self._file_mode = mode == "file"

class ProgressBarFlexibleWidget(ProgressBarWidget):
    r"""
    Progress bar widgets deriving from this abstract class indicate that they can expand/shrink freely.

    Thus, this class behaves as LaTeX's \hfill, or CSS flex: "1 1 auto".
    """

    def render(self, progress, width=None, status=None):
        """Render the widget."""
        raise NotImplementedError()

class _HiddenWidget(ProgressBarWidget):
    """Widget that doesn't render anything (for hidden progress bars)."""
    def render(self, progress, width=None, status=None):
        """Return empty render result."""
        return RenderResult()

class _ProgressBarCompoundWidget(ProgressBarWidget):
    """
    Container for other widgets.

    This widget is designed for internal use only! Its role is to lay out all contained widgets.
    """

    def __init__(self, widgets, title, file_mode):
        super(ProgressBarWidget, self).__init__()
        self._file_mode = file_mode
        self._width = min(self._get_terminal_size(), 100)
        self._encoding = (sys.stdout.encoding or "").lower()
        wlist = []
        for widget in (widgets or [title + ":", PBWBar(), PBWPercentage()]):
            if is_type(widget, str):
                widget = PBWString(widget)
            widget.set_mode("file" if file_mode else "tty")
            widget.set_encoding(self._encoding)
            wlist.append(widget)
        self._to_render = None  # Render this string on the next rendering cycle. Rarely used.
        self._widgets = tuple(wlist)
        self._widget_lengths = self._compute_widget_sizes()
        self._rendered = ""

    def render(self, progress, width=None, status=None):
        """Render the widget."""
        results = [widget.render(progress, width=self._widget_lengths[i], status=status)
                   for i, widget in enumerate(self._widgets)]

        if self._file_mode:
            res = ""
            for i, result in enumerate(results):
                res += result.rendered
                if result.length < self._widget_lengths[i] and progress < 1: break
                res += " " if i < len(results) - 1 else ""
            rendered_str = res[len(self._rendered):]
            self._rendered = res
        else:
            rendered_str = " ".join(r.rendered for r in results)
            if self._to_render:
                rendered_str = self._to_render + rendered_str
                self._to_render = None
        next_progress = min(r.next_progress for r in results)
        next_time = min(r.next_time for r in results)
        return RenderResult(rendered_str, next_progress=next_progress, next_time=next_time)

    def _compute_widget_sizes(self):
        """Initial rendering stage, done in order to compute widths of all widgets."""
        wl = [0] * len(self._widgets)
        flex_count = 0

        # First render all non-flexible widgets
        for i, widget in enumerate(self._widgets):
            if isinstance(widget, ProgressBarFlexibleWidget):
                flex_count += 1
            else:
                wl[i] = widget.render(1).length

        remaining_width = self._width - sum(wl)
        remaining_width -= len(self._widgets) - 1  # account for 1-space interval between widgets
        if remaining_width < 10 * flex_count:
            if self._file_mode:
                remaining_width = 10 * flex_count
            else:
                # The window is too small to accomodate the widget: try to split it into several lines, otherwise
                # switch to "file mode". If we don't do this, then rendering the widget will cause it to wrap, and
                # then when we use \r to go to the beginning of the line, only part of the widget will be overwritten,
                # which means we'll have many (possibly hundreds) of progress bar lines in the end.
                widget0 = self._widgets[0]
                if isinstance(widget0, PBWString) and remaining_width + widget0.render(0).length >= 10 * flex_count:
                    remaining_width += widget0.render(0).length + 1
                    self._to_render = widget0.render(0).rendered + "\n"
                    self._widgets = self._widgets[1:]
                if remaining_width < 10 * flex_count:
                    self._file_mode = True
                    remaining_width = 10 * flex_count

        remaining_width = max(remaining_width, 10 * flex_count)  # Ensure at least 10 chars per flexible widget

        for i, widget in enumerate(self._widgets):
            if isinstance(widget, ProgressBarFlexibleWidget):
                target_length = int(remaining_width / flex_count)
                result = widget.render(1, target_length)
                wl[i] = result.length
                remaining_width -= result.length
                flex_count -= 1

        return wl

    @staticmethod
    def _get_terminal_size():
        """Find current STDOUT's width, in characters."""
        # If output is not terminal but a regular file, assume 100 chars width
        if not sys.stdout.isatty():
            return 80

        # Otherwise, first try getting the dimensions from shell command `stty`:
        try:
            import subprocess
            ret = subprocess.check_output(["stty", "size"]).strip().split(" ")
            if len(ret) == 2:
                return int(ret[1])
        except:
            pass

        # Otherwise try using ioctl
        try:
            from termios import TIOCGWINSZ
            from fcntl import ioctl
            from struct import unpack
            res = unpack("hh", ioctl(sys.stdout, TIOCGWINSZ, b"1234"))
            return int(res[1])
        except:
            pass

        # Finally check the COLUMNS environment variable
        return int(os.environ.get("COLUMNS", 80))

#-----------------------------------------------------------------------------------------------------------------------
# Widgets (implementations)
#-----------------------------------------------------------------------------------------------------------------------

class PBWString(ProgressBarWidget):
    """Widget that represents a fixed string."""

    def __init__(self, text):
        """Initialize the widget."""
        super(ProgressBarWidget, self).__init__()
        self._str = "%s" % text

    def render(self, progress, width=None, status=None):
        """Render the widget."""
        return RenderResult(self._str)

class PBWBar(ProgressBarFlexibleWidget):
    """
    The widget that looks like an actual progress bar.

    It renders differently depending on whether the output supports Unicode or not:
        |███████   |  or
        [#######   ]
    """

    def __init__(self):
        """Initialize the widget."""
        super(ProgressBarWidget, self).__init__()
        self._bar_ends = None
        self._bar_symbols = None
        self._rendered = ""
        self.set_encoding(None)

    def render(self, progress, width=None, status=None):
        """Render the widget."""
        if width <= 3: return RenderResult()
        bar_width = width - 2  # Total width minus the bar ends

        n_chars = int(progress * bar_width + 0.001)
        endf, endl = self._bar_ends
        if self._file_mode:
            out = endf
            out += self._bar_symbols[-1] * n_chars
            out += endl if progress == 1 else ""
            if status:
                out += " (%s)" % status
            next_progress = (n_chars + 1) / bar_width
            rendered_len = len(out)
        else:
            frac_chars = int((progress * bar_width - n_chars) * len(self._bar_symbols))
            out = endf
            out += self._bar_symbols[-1] * n_chars
            out += self._bar_symbols[frac_chars - 1] if frac_chars > 0 else ""
            rendered_len = len(out)
            if status:
                out += colorama.Fore.RED + " (" + status + ")" + colorama.Style.RESET_ALL
                rendered_len += 3 + len(status)
            out += " " * (width - 1 - rendered_len)
            out += endl
            next_progress = (n_chars + (frac_chars + 1) / len(self._bar_symbols)) / bar_width
            rendered_len += max(0, width - 1 - rendered_len) + 1
        return RenderResult(rendered=out, length=rendered_len, next_progress=next_progress)

    def set_encoding(self, encoding):
        """Inform the widget about the encoding of the underlying character stream."""
        self._bar_ends = "[]"
        self._bar_symbols = "#"
        if not encoding: return
        s1 = "\u258F\u258E\u258D\u258C\u258B\u258A\u2589\u2588"
        s2 = "\u258C\u2588"
        s3 = "\u2588"
        if self._file_mode:
            s1 = s2 = None
        assert len(s3) == 1
        for s in (s1, s2, s3):
            if s is None: continue
            try:
                s.encode(encoding)
                self._bar_ends = "||"
                self._bar_symbols = s
                return
            except UnicodeEncodeError:
                pass
            except LookupError:
                print("Warning: unknown encoding %s" % encoding)

class PBWPercentage(ProgressBarWidget):
    """
    Simple percentage indicator.

    Renders like this: " 87%" (has constant width of 4 chars). Tries to re-render at every percentage point
    (so that you will see " 88%", " 89%", " 90%", etc).
    """

    def render(self, progress, width=None, status=None):
        """Render the widget."""
        current_pct = int(progress * 100 + 0.1)
        return RenderResult(rendered="%3d%%" % current_pct, next_progress=(current_pct + 1) / 100)

# Initialize colorama
colorama.init()


import sys
#sys.path.insert(1,"../../")
import h2o
#from tests import pyunit_utils
from random import randint
from random import uniform
from random import shuffle
import numpy as np

# This test is used to generate a dataframe that contains two columns, one integer and one double which
# should contain both positive, negative, zero.  Used to test sort float and the new TopN command.
def gen_data():
    floatA = []
    intA = []
    sizeMat = range(0,30)   # use to generate data of values 0, +/- 2^0 to +/1 2^64
    lowBoundF = -100000
    upperBoundF = -1*lowBoundF      # 2 million rows
    upperBoundL = pow(2,35)
    lowBoundL = upperBoundL-100000
    numZeros = 0
    numNans = 0   # generate Nans
    numInfs = 500
    numRep = 2    # number of times to repeat array
    csvFile = "/Users/wendycwong/temp/TopBottomNRep4.csv"
    fMult = 1.1

    fintA = []
    ffloatA = []
    for ind in range(0,1000):
        floatA = []
        intA = []
        genRandomData(intA,floatA, sizeMat)
        fintA.extend(intA)
        ffloatA.extend(floatA)

    shuffle(fintA)
    shuffle(ffloatA)
    bottom20FrameL = h2o.H2OFrame(python_obj=zip(fintA))
    bottom20FrameF = h2o.H2OFrame(python_obj=zip(ffloatA))
    h2o.download_csv(bottom20FrameL.cbind(bottom20FrameF), "/Users/wendycwong/temp/smallIntFloats.csv" )

    genStaticData(intA, floatA, upperBoundL, lowBoundF, upperBoundF, fMult)
    # save the correct sequence before shuffling for comparison purpose
    tempL = intA[0:int(round(len(intA)*0.2))]   # comes in decreasing value
    tempF = floatA[0:int(round(len(floatA)*0.2))]   # comes in decreasing value
    # save the correct sequence before shuffling for comparison purpose
    bottom20FrameL = h2o.H2OFrame(python_obj=zip(tempL))
    bottom20FrameF = h2o.H2OFrame(python_obj=zip(tempF))
    h2o.download_csv(bottom20FrameL.cbind(bottom20FrameF), "/Users/wendycwong/temp/Bottom20Per.csv" )

    tempL = intA[int(round(len(intA)*0.8)):len(intA)]
    tempL.sort()
    tempF = floatA[int(round(len(floatA)*0.8)):len(floatA)]
    tempF.sort()
    bottom20FrameL = h2o.H2OFrame(python_obj=zip(tempL))
    bottom20FrameF = h2o.H2OFrame(python_obj=zip(tempF))
    h2o.download_csv(bottom20FrameL.cbind(bottom20FrameF), "/Users/wendycwong/temp/Top20Per.csv" )

    # repeat the columns a few times to seriously test the algo with duplicated data.

    for val in range(0, numRep):
        intA.extend(intA)
        floatA.extend(floatA)

    shuffle(intA)   # randomly shuffle the indices
    shuffle(floatA) #

    intFrame = h2o.H2OFrame(python_obj=zip(intA))
    floatFrame = h2o.H2OFrame(python_obj=zip(floatA))
    h2o.download_csv(intFrame.cbind(floatFrame), csvFile)

def genDataFrame(sizeMat, lowBound, uppderBound, numRep, numZeros, numNans, numInfs):
    '''
    This function will generate an H2OFrame of two columns.  One column will be float and the other will
    be long.

    :param sizeMat: integer denoting size of bounds
    :param lowBound: lower bound
    :param uppderBound:
    :param trueRandom:
    :param numRep: number of times to repeat arrays in order to generate duplicated rows
    :param numZeros:
    :param numNans:
    :param numInfs:
    :return:
    '''
    if (numNans > 0):
        floatA = [float('NaN')]*numNans
        intA = [float('NaN')]*numNans

    if (numInfs > 0):
        floatA.extend([float('inf')]*numInfs)
        intA.extend([float('inf')]*numInfs)
        floatA.extend([-1.0*float('inf')]*numInfs)
        intA.extend([-1*float('inf')]*numInfs)

    # first generate the zeros into floatA and intA.  Multiply them with +/-1 to mess them up a little
    for index in range(numZeros):
        floatA.append(0.0*randint(-1,1))
        intA.append(0*randint(-1,1))

    # next generate +/- integers or various magnitude
    for rad in sizeMat:
        tempInt = pow(2,rad)
        tempIntN = pow(2,rad+1)
        intA.append(tempInt)
        intA.append(-1*tempInt)
        randInt = randint(tempInt, tempIntN)    # randomly generated integer, add both +/- values
        intA.append(randInt)
        intA.append(-1*randInt)
        intA.append(randint(tempInt, tempIntN))     # randomly generated integer
        intA.append(-1*(randint(tempInt, tempIntN)))    # randomly generated negative integer

        floatA.append(tempInt*1.0)
        floatA.append(-1.0*tempInt)
        tempD = uniform(tempInt, tempIntN)
        floatA.append(tempD)
        floatA.append(-1.0*tempD)
        floatA.append(uniform(tempInt, tempIntN))
        floatA.append(-1.0*uniform(tempInt, tempIntN))

def genRandomData(intA, floatA, sizeMat):
    # next generate +/- integers or various magnitude
    for rad in sizeMat:
        tempInt = pow(2,rad)
        tempIntN = pow(2,rad+1)
        intA.append(tempInt)
        intA.append(-1*tempInt)
        randInt = randint(tempInt, tempIntN)    # randomly generated integer, add both +/- values
        intA.append(randInt)
        intA.append(-1*randInt)
        intA.append(randint(tempInt, tempIntN))     # randomly generated integer
        intA.append(-1*(randint(tempInt, tempIntN)))    # randomly generated negative integer

        floatA.append(tempInt*1.0)
        floatA.append(-1.0*tempInt)
        tempD = uniform(tempInt, tempIntN)
        floatA.append(tempD)
        floatA.append(-1.0*tempD)
        floatA.append(uniform(tempInt, tempIntN))
        floatA.append(-1.0*uniform(tempInt, tempIntN))

def genStaticData(intA, floatA, upperBoundL, lowBoundF, upperBoundF, fMult):
    for val in range(lowBoundF, upperBoundF):
        floatA.append(val*fMult)
        intA.append(upperBoundL)
        upperBoundL=upperBoundL-1
    intA.reverse()

# This test will generate the merged data frame, the separate frame A and
# frame B, such that if you call merge on the separate frames, the answer is
# in the merged data frame.
def genMergedSeparaData(MergedRows, intUpper, intLow, doubleUpper, doubleLow, bProb):
    # first generate the single column that will be the merge key
    merged = h2o.create_frame(rows=MergedRows, cols=3, integer_fraction=1, integer_range=intUpper-intLow)
    print("Done, save with Flow")

def main(argv):
    h2o.init(strict_version_check=False)
    genMergedSeparaData(2000000000, pow(2,30), -1*pow(2,30), 1.0*pow(2,63), -1.0*pow(2,63), 0.8)

if __name__ == "__main__":
    main(sys.argv)
# if __name__ == "__main__":
#     pyunit_utils.standalone_test(gen_data)
# else:
#     gen_data()


# -*- encoding: utf-8 -*-
"""
H2O TargetEncoder.

:copyright: (c) 2016 H2O.ai
:license:   Apache License Version 2.0 (see LICENSE for details)
"""
from __future__ import absolute_import, division, print_function, unicode_literals

from h2o.expr import ExprNode
from h2o.frame import H2OFrame
from h2o.utils.typechecks import (assert_is_type)
from h2o import get_frame
import warnings

__all__ = ("TargetEncoder", )

class TargetEncoder(object):

    """
    Status: alpha version

    This is a main class that provides Python's API to the Java implementation of the target encoding.

    In general target encoding could be applied to three types of problems, namely:

         1) Binary classification (supported)
         2) Multi-class classification (not supported yet)
         3) Regression (not supported yet)

    Sample usage:

    >>> targetEncoder = TargetEncoder(x=te_columns, y=responseColumnName, blended_avg=True, inflection_point=3, smoothing=1)
    >>> targetEncoder.fit(trainFrame)
    >>> encodedTrain = targetEncoder.transform(frame=trainFrame, holdout_type="kfold", seed=1234, is_train_or_valid=True)
    >>> encodedValid = targetEncoder.transform(frame=validFrame, holdout_type="none", noise=0.0, is_train_or_valid=True)
    >>> encodedTest = targetEncoder.transform(frame=testFrame, holdout_type="none", noise=0.0, is_train_or_valid=False)
    """

    #-------------------------------------------------------------------------------------------------------------------
    # Construction
    #-------------------------------------------------------------------------------------------------------------------

    def __init__(self, x=None, y=None, fold_column='', blended_avg=True, inflection_point=3, smoothing=1, **kwargs):

        """
        Creates instance of the TargetEncoder class and setting parameters that will be used in both `train` and `transform` methods.

        :param List[str] or List[int] x: List of categorical column names or indices that we want apply target encoding to.

        :param str or int y: the name or column index of the response variable in the data.
        :param str or int fold_column: the name or column index of the fold column in the data.
        :param boolean blending_avg: (deprecated) whether to perform blended average. Defaults to TRUE.
        :param boolean blended_avg: whether to perform blended average. Defaults to TRUE.
        :param double inflection_point: parameter for blending. Used to calculate `lambda`. Determines half of the minimal sample size
            for which we completely trust the estimate based on the sample in the particular level of categorical variable.
        :param double smoothing: parameter for blending. Used to calculate `lambda`. Controls the rate of transition between
            the particular level's posterior probability and the prior probability. For smoothing values approaching infinity it becomes a hard
            threshold between the posterior and the prior probability.

        """

        if(type(x) == str or type(x) == int):
            x = [x]
        self._teColumns = x
        self._responseColumnName = y
        self._foldColumnName = fold_column
        if 'blending_avg' in kwargs:
            warnings.warn("Parameter blending_avg is deprecated; use blended_avg instead", category=DeprecationWarning, stacklevel=2)
            self._blending = kwargs.get('blending_avg')
        else:
            self._blending = blended_avg

        self._inflectionPoint = inflection_point
        self._smoothing = smoothing

    def fit(self, frame = None):
        """
        Returns encoding map as an object that maps 'column_name' -> 'frame_with_encoding_map_for_this_column_name'

        :param frame frame: An H2OFrame object with which to create the target encoding map
        """
        self._teColumns = list(map(lambda i: frame.names[i], self._teColumns)) if all(isinstance(n, int) for n in self._teColumns) else self._teColumns
        self._responseColumnName = frame.names[self._responseColumnName] if isinstance(self._responseColumnName, int) else self._responseColumnName
        self._foldColumnName = frame.names[self._foldColumnName] if isinstance(self._foldColumnName, int) else self._foldColumnName

        self._encodingMap = ExprNode("target.encoder.fit", frame, self._teColumns, self._responseColumnName,
                                     self._foldColumnName)._eager_map_frame()

        return self._encodingMap

    def transform(self, frame=None, holdout_type=None, noise=-1, seed=-1):
        """
        Apply transformation to `te_columns` based on the encoding maps generated during `TargetEncoder.fit()` call.
        You must not pass encodings manually from `.fit()` method because they are being stored internally
        after `.fit()' had been called.

        :param frame frame: to which frame we are applying target encoding transformations.
        :param str holdout_type: Supported options:

                1) "kfold" - encodings for a fold are generated based on out-of-fold data.
                2) "loo" - leave one out. Current row's response value is subtracted from the pre-calculated per-level frequencies.
                3) "none" - we do not holdout anything. Using whole frame for training

        :param float noise: the amount of random noise added to the target encoding.  This helps prevent overfitting. Defaults to 0.01 * range of y.
        :param int seed: a random seed used to generate draws from the uniform distribution for random noise. Defaults to -1.
        """
        assert_is_type(holdout_type, "kfold", "loo", "none")

        # We need to make sure that frames are being sent in the same order
        assert self._encodingMap.map_keys['string'] == self._teColumns
        encodingMapKeys = self._encodingMap.map_keys['string']
        encodingMapFramesKeys = list(map(lambda x: x['key']['name'], self._encodingMap.frames))
        return H2OFrame._expr(expr=ExprNode("target.encoder.transform", encodingMapKeys, encodingMapFramesKeys, frame, self._teColumns, holdout_type,
                                            self._responseColumnName, self._foldColumnName,
                                            self._blending, self._inflectionPoint, self._smoothing,
                                            noise, seed))

    def encoding_map_frames(self):
        return list(map(lambda x: get_frame(x['key']['name']), self._encodingMap.frames))