from bokeh import __version__
__version__

import numpy as np
from bokeh.plotting import figure
from bokeh.io import output_notebook, show
from bokeh.models import Row

output_notebook()

from numba import autojit

# These functions generate the Mandelbrot set image. Don't worry if
# you are not familiar with them. The import thing is just to know
# that they create a 2D array of numbers that we can colormap.

@autojit
def mandel(x, y, max_iters):
    """
    Given the real and imaginary parts of a complex number,
    determine if it is a candidate for membership in the Mandelbrot
    set given a fixed number of iterations.
    """
    c = complex(x, y)
    z = 0.0j
    for i in range(max_iters):
        z = z*z + c
        if (z.real*z.real + z.imag*z.imag) >= 4:
            return i
    return max_iters

@autojit
def create_fractal(min_x, max_x, min_y, max_y, image, iters):
    height = image.shape[0]
    width = image.shape[1]

    pixel_size_x = (max_x - min_x) / width
    pixel_size_y = (max_y - min_y) / height

    for x in range(width):
        real = min_x + x * pixel_size_x
        for y in range(height):
            imag = min_y + y * pixel_size_y
            color = mandel(real, imag, iters)
            image[y, x] = color

# Define the bounding coordinates to generate the Mandelbrot image in. You
# can play around with these.
min_x = -2.0
max_x = 1.0
min_y = -1.0
max_y = 1.0

# Use the functions above to create a scalar image (2D array of numbers)
img = np.zeros((1024, 1536), dtype = np.uint8)
create_fractal(min_x, max_x, min_y, max_y, img, 20)

# create a figure, setting the x and y ranges to the appropriate data bounds
p1 = figure(title="Mandelbrot", plot_width=int(900/2), plot_height=int(600/2),
            x_range = [min_x, max_x], y_range = [min_y, max_y])

# Fill in the missing parameters to use the `image` renderer to
# display the Mandelbrot image color mapped with the palette 'Spectral11'
#
# NOTE: the `image` renderer can display many images at once, so it takes
# **lists** of images, coordinates, and palettes. Remember to supply sequences
# for these parameters, even if you are just supplying one.
p1.image(image=[img],             # image data
         x=[min_x],               # lower left x coord
         y=[min_y],               # lower left y coord
         dw=[max_x-min_x],        # *data space* width of image
         dh=[max_y-min_y],        # *data space* height of image
         palette="Spectral11",    # palette name
)

# create a new figure
p2 = figure(title="RGBA image",plot_width=int(900/2), plot_height=int(600/2),
            x_range = [0,10], y_range = [0,10], )

# We can also use the `image_rgba` renderer to display RGBA images that
# we have color mapped ourselves.
N = 20
img = np.empty((N,N), dtype=np.uint32)
view = img.view(dtype=np.uint8).reshape((N, N, 4))
for i in range(N):
    for j in range(N):
        view[i, j, 0] = int(i/N*255) # red
        view[i, j, 1] = 158          # green
        view[i, j, 2] = int(j/N*255) # blue
        view[i, j, 3] = 255          # alpha

# Use `image_rgba` to display the image above. 
p2.image_rgba(image=[img], x=[0], y=[0], dw=[10], dh=[10])

# show the plots in a Row
show(Row(p1, p2))

from bokeh.models import CustomJS, ColumnDataSource, Slider, WidgetBox, Row, Column

fill_alpha = 0.5
init_color_value = 127
color = "rgb({0}, {0}, {0})".format(init_color_value)

source = ColumnDataSource(data=dict(fill_alpha=[fill_alpha], color=[color]))

callback = CustomJS(
    args=dict(source=source), 
    code="""var data = source.get('data');
    var r = r_slider.get('value');
    var g = g_slider.get('value');
    var b = b_slider.get('value');
    data.color = ["rgb(" + r + ", " + g + ", " + b + ")"];
    data.fill_alpha = [a_slider.get('value')];
    source.trigger('change');
    """
)

plot = figure(plot_width=335, plot_height=335, title="RGBA Color Plot", toolbar_location=None, tools="", outline_line_color="#FFFFFF")
plot.grid.grid_line_color=None
plot.axis.visible = False

plot.rect(x=[1,], y=[1,], width=5, height=15, source=source, color="color", fill_alpha="fill_alpha")

r_slider = Slider(start=0, end=255, value=init_color_value, step=1, title="R", callback=callback)
g_slider = Slider(start=0, end=255, value=init_color_value, step=1, title="G", callback=callback)
b_slider = Slider(start=0, end=255, value=init_color_value, step=1, title="B", callback=callback)
a_slider = Slider(start=0, end=1, value=fill_alpha, step=0.01, title="A", callback=callback)

callback.args['r_slider'] = r_slider
callback.args['g_slider'] = g_slider
callback.args['b_slider'] = b_slider
callback.args['a_slider'] = a_slider

layout = Row(
    WidgetBox(r_slider, g_slider, b_slider, a_slider, width=300), plot
)

show(layout)