import numpy as np

from bokeh.plotting import figure
from bokeh.models import ColumnDataSource, ColorBar, LinearColorMapper

from buki import Object


class Plot2D(Object):

    def __init__(self, name=None):
        data = {
            "image": [np.zeros((1, 1))]
        }

        self.source = source = ColumnDataSource(data=data)

        cmap = LinearColorMapper(palette="Viridis256")
        cbar = ColorBar(color_mapper=cmap)

        fig = figure(name=name, match_aspect=True)
        fig.image(source=source, x=0, y=0, dw=1, dh=1, color_mapper=cmap)

        fig.x_range.range_padding = 0
        fig.y_range.range_padding = 0

        fig.add_layout(cbar, "right")
        super().__init__(fig)


    def set(self, times, values):
        image = values[-1]
        self.source.data.update(image=[image])