import numpy as np
from numba import autojit, jit, double
%pylab inline

"""
Explanation: This example assumes the notebook server has been called with ipython notebook --pylab inline and the trunk version of numba at Github.
End of explanation
"""


def sum(arr):
    M, N = arr.shape
    sum = 0.0
    for i in range(M):
        for j in range(N):
            sum += arr[i,j]
    return sum
fastsum = jit('f8(f8[:,:])')(sum)
flexsum = autojit(sum)

arr2d = np.arange(600,dtype=float).reshape(20,30)
print(sum(arr2d))
print(fastsum(arr2d))
print(flexsum(arr2d))
print(flexsum(arr2d.astype(int)))

%timeit sum(arr2d)

%timeit fastsum(arr2d)

416 / .921# speedup

%timeit arr2d.sum() 

7.86 / .921 # even provides a speedup over general-purpose NumPy sum

"""
Explanation: Numba provides two major decorators:  jit and autojit.  
The jit decorator returns a compiled version of the function using the input types and the output types of the function.   You can specify the type using out_type(in_type, ...) syntax.  Array inputs can be specified using [:,:] appended to the type.  
The autojit decorator does not require you to specify any types.  It watches for what types you call the function with and infers the type of the return.  If there is a previously compiled version of the code available it uses it, if not it generates machine code for the function and then executes that code.
End of explanation
"""


@jit('void(f8[:,:],f8[:,:],f8[:,:])')
def filter(image, filt, output):
    M, N = image.shape
    m, n = filt.shape
    for i in range(m//2, M-m//2):
        for j in range(n//2, N-n//2):
            result = 0.0
            for k in range(m):
                for l in range(n):
                    result += image[i+k-m//2,j+l-n//2]*filt[k, l]
            output[i,j] = result

import urllib
bytes = urllib.urlopen('http://www.cs.tut.fi/~foi/SA-DCT/original/image_Lake512.png').read()            
from matplotlib.pyplot import imread
import cStringIO
image = imread(cStringIO.StringIO(bytes)).astype('double')
import time
filt = np.ones((15,15),dtype='double')
filt /= filt.sum()
output = image.copy()
filter(image, filt, output)
gray()
imshow(output)
start = time.time()
filter(image[:100,:100], filt, output[:100,:100])
fast = time.time() - start
start = time.time()
filter.py_func(image[:100,:100], filt, output[:100,:100])
slow = time.time() - start
print("Python: %f s; Numba: %f ms; Speed up is %f" % (slow, fast*1000, slow / fast))

imshow(image)
gray()

"""
Explanation: The speed-up is even more pronounced the more inner loops in the code.   Here is an image processing example:
End of explanation
"""


@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.
    """
    i = 0
    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 255

@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

    return image

image = np.zeros((500, 750), dtype=np.uint8)
imshow(create_fractal(-2.0, 1.0, -1.0, 1.0, image, 20))
jet()

%timeit create_fractal(-2.0, 1.0, -1.0, 1.0, image, 20)

%timeit create_fractal.py_func(-2.0, 1.0, -1.0, 1.0, image, 20)

2.14/16.3e-3 

"""
Explanation: You can call Numba-created functions from other Numba-created functions and get even more amazing speed-ups.
End of explanation
"""


from numba import double, autojit

class MyClass(object):
    def mymethod(self, arg):
        return arg * 2

@autojit(locals=dict(mydouble=double)) # specify types for local variables
def call_method(obj):
    print(obj.mymethod("hello"))        # object result
    mydouble = obj.mymethod(10.2) # native double
    print(mydouble * 2)                      # native multiplication

call_method(MyClass())

"""
Explanation: Numba works very well for numerical calculations and infers types for variables.  You can over-ride this inference by passing in a locals dictionary to the autojit decorator.   Notice how the code below shows both Python object manipulation and native manipulation
End of explanation
"""


@autojit
def complex_support(real, imag):
    c = complex(real, imag)
    return (c ** 2).conjugate()

c = 2.0 + 4.0j
complex_support(c.real, c.imag), (c**2).conjugate()

"""
Explanation: Basic complex support is available as well.  Some functions are still being implemented, however.
End of explanation
"""


from numba import struct, jit, double
import numpy as np

record_type = struct([('x', double), ('y', double)])
record_dtype = record_type.get_dtype()
a = np.array([(1.0, 2.0), (3.0, 4.0)], dtype=record_dtype)

@jit(argtypes=[record_type[:]])
def hypot(data):
    # return types of numpy functions are inferred
    result = np.empty_like(data, dtype=np.float64)
    # notice access to structure elements 'x' and 'y' via attribute access
    # You can also index by field name or field index:
    #       data[i].x == data[i]['x'] == data[i][0]
    for i in range(data.shape[0]):
        result[i] = np.sqrt(data[i].x * data[i].x + data[i].y * data[i].y)

    return result

print(hypot(a))

# Notice inferred return type
print(hypot.signature)
# Notice native sqrt calls and for.body direct access to memory...
print(hypot.lfunc)

print(hypot.signature) # inspect function signature, note inferred return type

[line for line in str(hypot.lfunc).splitlines() if 'sqrt' in line] # note native math calls

"""
Explanation: We can even create a function that takes a structured array as input.
End of explanation
"""


