import pyopencl as cl
import numpy as np
from PIL import Image
from decimal import Decimal

def mandel(ctx, x, y, zoom, max_iter=1000, iter_steps=1, width=500, height=500, use_double=False):
    mf = cl.mem_flags
    cl_queue = cl.CommandQueue(ctx)
    # build program
    code = """
    #if real_t == double
        #pragma OPENCL EXTENSION cl_khr_fp64 : enable
    #endif
    kernel void mandel(
        __global real_t2 *coords,
        __global uint *ids_map,
        __global uint *output,
        uint max_iter,
        uint start_iter
    ){
        uint id = get_global_id(0);
        uint rid = ids_map[id];
        real_t2 my_coords = coords[rid];
        real_t x = my_coords.x;
        real_t y = my_coords.y;
        uint iter = 0;
        for(iter=start_iter; iter<max_iter; ++iter){
            if(x*x + y*y > 4.0f){
                break;
            }
            real_t xtemp = x*x - y*y + my_coords.x;
            y = 2*x*y + my_coords.y;
            x = xtemp;
        }
        // copy the current x,y pair back
        my_coords.x = x;
        my_coords.y = y;
        coords[rid] = my_coords;
        output[id] = iter;
    }
    """
    _cltype, _nptype = ("double",np.float64) if use_double else ("float", np.float32)
    prg = cl.Program(ctx, code).build("-cl-opt-disable -D real_t=%s -D real_t2=%s2" % (_cltype, _cltype))

    # Calculate the "viewport".
    x0 = x - ((Decimal(3) * zoom)/Decimal(2.))
    y0 = y - ((Decimal(2) * zoom)/Decimal(2.))
    x1 = x + ((Decimal(3) * zoom)/Decimal(2.))
    y1 = y + ((Decimal(2) * zoom)/Decimal(2.))

    # Create index map in x,y pairs
    xx = np.arange(_nptype(x0), _nptype(x1), _nptype((x1-x0)/Decimal(width)), dtype=_nptype)
    yy = np.arange(_nptype(y0), _nptype(y1), _nptype((y1-y0)/Decimal(height)), dtype=_nptype)
    coord_map = np.dstack(np.meshgrid(xx, yy)).astype(_nptype).flatten()
    buffer_coords_cl = cl.Buffer(ctx, mf.COPY_HOST_PTR, hostbuf=coord_map)
    # Create input and output buffer
    indices = np.arange(0, width*height, 1, np.uint32)
    buffer_in_cl = cl.Buffer(ctx, mf.READ_ONLY, size=indices.nbytes)
    buffer_out = np.zeros(width*height, dtype=np.uint32) # This will contain the iteration values of that run
    buffer_out_cl = cl.Buffer(ctx, mf.WRITE_ONLY, size=buffer_out.nbytes)
    # 2D Buffer to collect the iterations needed per pixel
    iter_map = np.zeros(width*height, dtype=np.uint32)

    start_max_iter = 0
    to_do = indices.size
    steps_size = int(max_iter / float(iter_steps))
    while to_do > 0 and start_max_iter < max_iter:
        end_max_iter = min(max_iter, start_max_iter + steps_size )
        print "Iterations from iteration %i to %i for %i numbers" % (start_max_iter, end_max_iter, to_do)
        # copy indices to device
        cl.enqueue_copy(cl_queue, buffer_in_cl, indices[:to_do]).wait()
        # and finally call the ocl function
        prg.mandel(cl_queue, (to_do,), None,
            buffer_coords_cl,
            buffer_in_cl,
            buffer_out_cl,
            np.uint32(end_max_iter),
            np.uint32(start_max_iter)
        ).wait()
        # Copy the output back
        cl.enqueue_copy(cl_queue, buffer_out, buffer_out_cl).wait()
        # Get indices of "found" escapes
        done = np.where(buffer_out[:to_do] < end_max_iter)[0]
        done_index = indices[done]
        # Write needed iteration to out buffer
        iter_map[done_index] = buffer_out[done]
        # Write non escapes back to indices for the next loop
        undone = np.where(buffer_out[:to_do] >= end_max_iter)[0]
        indices[:undone.size] = indices[undone]
        to_do = undone.size
        start_max_iter = end_max_iter
        print "%i done. %i unknown" % (done.size, undone.size)

    # simple coloring by modulo 255 on the iter_map
    return (iter_map % 255).astype(np.uint8).reshape((height, width))


if __name__ == '__main__':
    ctx = cl.create_some_context(interactive=True)
    img = mandel(ctx,
          x=Decimal("-0.7546546453361122021732941811"),
          y=Decimal("0.05020518634419688663435986387"),
          zoom=Decimal("0.0002046859427855630601247281079"),
          max_iter=2000,
          iter_steps=10,
          width=500,
          height=400,
          use_double=False
    )
    Image.fromarray(img).show()