from pywr.recorders import Recorder
import time
import logging
logger = logging.getLogger(__name__)


class ProgressRecorder(Recorder):
    """Simple text-based progress notifications

    Parameters
    ----------
    print_func : callable or None
        The function to call when updating progress. The function is given a single str argument that
        contains the progress message. Defaults to `logger.info`.
    """
    def __init__(self, *args, **kwargs):
        print_func = kwargs.pop('print_func', None)
        super(ProgressRecorder, self).__init__(*args, **kwargs)
        if print_func is None:
            print_func = logger.info
        self.print_func = print_func
        self.last_progress = None
        self.last_timestep = None
        self.t0 = None
        self.combinations = None

    def reset(self):
        self.last_progress = -1
        self.last_timestep = 0
        self.t0 = time.time()
        self.combinations = len(self.model.scenarios.combinations)

    def after(self):
        total_timesteps = len(self.model.timestepper)
        timestep = self.model.timestepper.current.index
        progress = int((timestep+1) / total_timesteps * 100)
        if progress > self.last_progress:
            self.last_progress = progress
            if progress >= 1:
                time_taken = time.time() - self.t0
                try:
                    speed = ((timestep-self.last_timestep)*self.combinations) / time_taken
                except ZeroDivisionError:
                    speed = float("inf")
                self.last_timestep = timestep
                self.update_progress(progress, speed)
            else:
                self.last_timestep = timestep
                self.update_progress(progress, None)
            self.t0 = time.time()
            self.last_timestep = timestep

    def update_progress(self, progress, speed=None):
        if speed is not None:
            self.print_func("Completed {}%, {:.0f} steps/second".format(progress, speed))
        else:
            self.print_func("Completed {}%".format(progress))

class JupyterProgressRecorder(ProgressRecorder):
    """Graphical progress bar for use in Jupyter notebooks"""
    def __init__(self, *args, **kwargs):
        super(JupyterProgressRecorder, self).__init__(*args, **kwargs)

    def reset(self):
        from ipywidgets import FloatProgress, HBox, Label, Layout
        from IPython.display import display
        super(JupyterProgressRecorder, self).reset()
        self.progress_bar = FloatProgress(min=0, max=100, description='Running:')
        self.label = Label("", layout=Layout(width='100%'))
        self.box = HBox([self.progress_bar, self.label])
        display(self.box)

    def update_progress(self, progress, speed):
        self.progress_bar.value = progress
        if speed is not None:
            self.label.value = "{:.0f} steps/second".format(speed)
        else:
            self.label.value = ""

    def finish(self):
        super(JupyterProgressRecorder, self).finish()
        if self.progress_bar.value >= 100.0:
            self.progress_bar.bar_style = "success"
        else:
            self.progress_bar.bar_style = "danger"

