%%time
import numpy as np
from numba import cuda
from numba import *

from matplotlib import pyplot as plt
%matplotlib inline
plt.rcParams['figure.figsize'] = [15, 5]
plt.rcParams['figure.dpi'] = 100

@jit
def mandel(real, imag, max_iters):
    cr = real
    ci = imag
    for n in range(max_iters):
        real2 = real*real
        imag2 = imag*imag
        if real2 + imag2 > 4.0:
            return n
        imag = 2*real*imag+ci
        real = real2-imag2+cr
    return max_iters

mandel_gpu = cuda.jit(uint32(float32, float32, uint32), device=True)(mandel)

@cuda.jit((float32, float32, float32, float32, uint32[:,:], uint32))
def mandel_kernel(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

  startX, startY = cuda.grid(2)
  gridX = cuda.gridDim.x * cuda.blockDim.x;
  gridY = cuda.gridDim.y * cuda.blockDim.y;

  for x in range(startX, width, gridX):
    real = min_x + x * pixel_size_x
    for y in range(startY, height, gridY):
      imag = min_y + y * pixel_size_y 
      image[y][x] = mandel_gpu(real, imag, iters)
        
gimage = np.zeros((1000, 1000), dtype = np.uint32)
blockdim = (32, 8)
griddim = (32,16)

d_image = cuda.to_device(gimage)
mandel_kernel[griddim, blockdim](-2.0, 1.0, -1.0, 1.0, d_image, 20) 
d_image.copy_to_host(gimage)
plt.subplot(1, 2, 1).imshow(gimage)

d_image = cuda.to_device(gimage)
mandel_kernel[griddim, blockdim](-0.74877,-0.74872,0.06505,0.06510, d_image, 2048) 
d_image.copy_to_host(gimage)
plt.subplot(1, 2, 2).imshow(gimage)