from __future__ import print_function, division

from timeit import default_timer as timer

from bokeh.plotting import figure, show, output_notebook
from bokeh.models import GlyphRenderer, LinearColorMapper
from bokeh.io import push_notebook
from numba import jit, njit

from ipywidgets import interact
import numpy as np
import scipy.misc

output_notebook()

"""
Explanation: Interactive Image Processing with Numba and Bokeh
This demo shows off how interactive image processing can be done in the notebook, using Numba for numerics, Bokeh for plotting, and Ipython interactors for widgets. The demo runs entirely inside the Ipython notebook, with no Bokeh server required.
Numba must be installed in order to run this demo. To run, click on, Cell-&gt;Run All in the top menu, then scroll down to individual examples and play around with their controls.
End of explanation
"""


# smaller image
img_blur = (scipy.misc.ascent()[::-1,:]/255.0)[:250, :250].copy(order='C')

palette = ['#%02x%02x%02x' %(i,i,i) for i in range(256)]
width, height = img_blur.shape
p_blur = figure(x_range=(0, width), y_range=(0, height))
r_blur = p_blur.image(image=[img_blur], x=[0], y=[0], dw=[width], dh=[height], palette=palette, name='blur')

@njit
def blur(outimg, img, amt):
    iw, ih = img.shape
    for i in range(amt, iw-amt):
        for j in range(amt, ih-amt):
            px = 0.
            for w in range(-amt//2, amt//2):
                for h in range(-amt//2, amt//2):
                    px += img[i+w, j+h]
            outimg[i, j]= px/(amt*amt)

def update(i=0):
    level = 2*i + 1
    
    out = img_blur.copy()
    
    ts = timer()
    blur(out, img_blur, level)
    te = timer()
    print('blur takes:', te - ts)
    
    renderer = p_blur.select(dict(name="blur", type=GlyphRenderer))
    r_blur.data_source.data['image'] = [out]
    push_notebook(handle=t_blur)

t_blur = show(p_blur, notebook_handle=True)

interact(update, i=(0, 10))

"""
Explanation: Gaussian Blur
This first section demonstrates performing a simple Gaussian blur on an image. It presents the image, as well as a slider that controls how much blur is applied. Numba is used to compile the python blur kernel, which is invoked when the user modifies the slider. 
Note: This simple example does not handle the edge case, so the edge of the image will remain unblurred as the slider is increased.
End of explanation
"""


@jit
def getitem(img, x, y):
    w, h = img.shape
    if x >= w:
        x = w - 1 - (x - w)
    if y >= h:
        y = h - 1 - (y - h)
    return img[x, y]
 
def filter_factory(kernel):
    ksum = np.sum(kernel)
    if ksum == 0:
        ksum = 1
    k9 = kernel / ksum
 
    @jit
    def kernel_apply(img, out, x, y):
        tmp = 0
        for i in range(3):
            for j in range(3):
                tmp += img[x+i-1, y+j-1] * k9[i, j]
        out[x, y] = tmp
 
    @jit
    def kernel_apply_edge(img, out, x, y):
        tmp = 0
        for i in range(3):
            for j in range(3):
                tmp += getitem(img, x+i-1, y+j-1) * k9[i, j]
        out[x, y] = tmp
 
    @jit
    def kernel_k9(img, out):
        # Loop through all internals
        for x in range(1, img.shape[0] -1):
            for y in range(1, img.shape[1] -1):
                kernel_apply(img, out, x, y)
 
        # Loop through all the edges
        for x in range(img.shape[0]):
            kernel_apply_edge(img, out, x, 0)
            kernel_apply_edge(img, out, x, img.shape[1] - 1)
 
        for y in range(img.shape[1]):
            kernel_apply_edge(img, out, 0, y)
            kernel_apply_edge(img, out, img.shape[0] - 1, y)
 
    return kernel_k9

average = np.array([
    [1, 1, 1],
    [1, 1, 1],
    [1, 1, 1],
], dtype=np.float32)

sharpen = np.array([
    [-1, -1, -1],
    [-1, 12, -1],
    [-1, -1, -1],
], dtype=np.float32)

edge = np.array([
    [ 0, -1,  0],
    [-1,  4, -1],
    [ 0, -1,  0],
], dtype=np.float32)

edge_h = np.array([
    [ 0,  0,  0],
    [-1,  2, -1],
    [ 0,  0,  0],
], dtype=np.float32)

edge_v = np.array([
    [0, -1, 0],
    [0,  2, 0],
    [0, -1, 0],
], dtype=np.float32)

gradient_h = np.array([
    [-1, -1, -1],
    [ 0,  0,  0],
    [ 1,  1,  1],
], dtype=np.float32)

gradient_v = np.array([
    [-1, 0, 1],
    [-1, 0, 1],
    [-1, 0, 1],
], dtype=np.float32)

sobol_h = np.array([
    [ 1,  2,  1],
    [ 0,  0,  0],
    [-1, -2, -1],
], dtype=np.float32)

sobol_v = np.array([
    [-1, 0, 1],
    [-2, 0, 2],
    [-1, 0, 1],
], dtype=np.float32)
 
emboss = np.array([    
    [-2, -1, 0],
    [-1,  1, 1],
    [ 0,  1, 2],
], dtype=np.float32)

kernels = {
    "average"               : filter_factory(average),
    "sharpen"               : filter_factory(sharpen),
    "edge (both)"           : filter_factory(edge),
    "edge (horizontal)"     : filter_factory(edge_h),
    "edge (vertical)"       : filter_factory(edge_v),
    "gradient (horizontal)" : filter_factory(gradient_h),
    "gradient (vertical)"   : filter_factory(gradient_v),
    "sobol (horizontal)"    : filter_factory(sobol_h),
    "sobol (vertical)"      : filter_factory(sobol_v),
    "emboss"                : filter_factory(emboss),
}

images = {
    "ascent" : np.copy(scipy.misc.ascent().astype(np.float32)[::-1, :]),
    "face"   : np.copy(scipy.misc.face(gray=True).astype(np.float32)[::-1, :]),
}

palette = ['#%02x%02x%02x' %(i,i,i) for i in range(256)]
cm = LinearColorMapper(palette=palette, low=0, high=256)
width, height = images['ascent'].shape
p_kernel = figure(x_range=(0, width), y_range=(0, height))
r_kernel = p_kernel.image(image=[images['ascent']], x=[0], y=[0], dw=[width], dh=[height], color_mapper=cm, name="kernel")

def update(image="ascent", kernel_name="none", scale=100, bias=0):
    global _last_kname
    global _last_out
    
    img_kernel = images.get(image)

    kernel = kernels.get(kernel_name, None)
    if kernel == None:
        out = np.copy(img_kernel)

    else:
        out = np.zeros_like(img_kernel)

        ts = timer()
        kernel(img_kernel, out)
        te = timer()
        print('kernel takes:', te - ts)

    out *= scale / 100.0
    out += bias

    r_kernel.data_source.data['image'] = [out]
    push_notebook(handle=t_kernel)

t_kernel = show(p_kernel, notebook_handle=True)

knames = ["none"] + sorted(kernels.keys())
interact(update, image=["ascent" ,"face"], kernel_name=knames, scale=(10, 100, 10), bias=(0, 255))

"""
Explanation: 3x3 Image Kernels
Many image processing filters can be expressed as 3x3 matrices. This more sophisticated example demonstrates how numba can be used to compile kernels for arbitrary 3x3 kernels, and then provides serveral predefined kernels for the user to experiment with. 
The UI presents the image to process (along with a dropdown to select a different image) as well as a dropdown that lets the user select which kernel to apply. Additioanlly there are sliders the permit adjustment to the bias and scale of the final greyscale image. 
Note: Right now, adjusting the scale and bias are not as efficient as possible, because the update function always also applies the kernel (even if it has not changed). A better implementation might have a class that keeps track of the current kernal and output image so that bias and scale can be applied by themselves.
End of explanation
"""


@njit
def wavelet_decomposition(img, tmp):
    """
    Perform inplace wavelet decomposition on `img` with `tmp` as
    a temporarily buffer.

    This is a very simple wavelet for demonstration
    """
    w, h = img.shape
    halfwidth, halfheight = w//2, h//2
 
    lefthalf, righthalf = tmp[:halfwidth, :], tmp[halfwidth:, :]
 
    # Along first dimension
    for x in range(halfwidth):
        for y in range(h):
            lefthalf[x, y] = (img[2 * x, y] + img[2 * x + 1, y]) / 2
            righthalf[x, y] = img[2 * x, y] - img[2 * x + 1, y]
 
    # Swap buffer
    img, tmp = tmp, img
    tophalf, bottomhalf = tmp[:, :halfheight], tmp[:, halfheight:]
 
    # Along second dimension
    for y in range(halfheight):
        for x in range(w):
            tophalf[x, y] = (img[x, 2 * y] + img[x, 2 * y + 1]) / 2
            bottomhalf[x, y] = img[x, 2 * y] - img[x, 2 * y + 1]
 
    return halfwidth, halfheight

img_wavelet = np.copy(scipy.misc.face(gray=True)[::-1, :])

palette = ['#%02x%02x%02x' %(i,i,i) for i in range(256)]
width, height = img_wavelet.shape
p_wavelet = figure(x_range=(0, width), y_range=(0, height))
r_wavelet = p_wavelet.image(image=[img_wavelet], x=[0], y=[0], dw=[width], dh=[height], palette=palette, name="wavelet")

def update(level=0):

    out = np.copy(img_wavelet)
    tmp = np.zeros_like(img_wavelet)

    ts = timer()
    hw, hh = img_wavelet.shape
    while level > 0 and hw > 1 and hh > 1:
        hw, hh = wavelet_decomposition(out[:hw, :hh], tmp[:hw, :hh])
        level -= 1
    te = timer()
    print('wavelet takes:', te - ts)

    r_wavelet.data_source.data['image'] = [out]
    push_notebook(handle=t_wavelet)

t_wavelet = show(p_wavelet, notebook_handle=True)

interact(update, level=(0, 7))

"""
Explanation: Wavelet Decomposition
This last example demostrates a Haar wavelet decomposition using a Numba-compiled function. Play around with the slider to see differnet levels of decomposition of the image.
End of explanation
"""


