from bokeh.embed import server_document
from bokeh.layouts import column
from bokeh.models import ColumnDataSource, Button, CustomJS
from bokeh.plotting import figure, curdoc
from bokeh.server.server import Server
from bokeh.themes import Theme

import socket
import numpy as np
from pprint import pprint

plot = figure(x_range=(0,600), y_range=(0,600))

N = 1000
x = np.linspace(0, 10, N, dtype=np.float32)
y = np.linspace(0, 10, N, dtype=np.float32)
xx, yy = np.meshgrid(x, y)
d = np.sin(xx)*np.cos(yy)

img_source=ColumnDataSource(data=dict(image=[ d[  0:100,   0:100] ],
                                      x=[0],
                                      y=[0],
                                      dw=[100],
                                      dh=[100]))
plot.image(source=img_source, 
           image='image', x='x', y='y', dw='dw', dh='dh',
           name="tile", dilate=True)
i=1

def callback():
    global i
    nimg=len(img_source.data["x"])
    print "callback. %d images" %  nimg
    if i<10:
        stream_dict=dict(image=[ d[ i*100:(i+1)*100, i*100:(i+1)*100 ] ],
                         x=[i*100],
                         y=[i*100],
                         dw=[100],
                         dh=[100])
        img_source.stream(stream_dict)
        i+=1
    # Doing this makes it work, but it's slow. Is it resending the whole data?
    # img_source.trigger('data', img_source.data, img_source.data)
button = Button(label="Add a plot")
button.on_click(callback)

# This doesn't quite work either
# js_callback=CustomJS(args=dict(img_source=img_source), code="""
# var nimg=img_source.data.image.length;
# img_source._shapes.image[nimg-1]=[100,100];
# img_source.change.emit();
# """)
# img_source.js_on_change('streaming', js_callback)

curdoc().add_root(column(button, plot))