import pyopencl as cl
import numpy
import numpy.linalg as la
from scipy.misc import pilutil

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)

def dostuff(rays, map):
    ray_floats, ray_ints = rays
    num_rays = len(ray_floats)
    
    ray_floats = ray_floats.ravel()
    ray_ints = ray_ints.ravel()
    print ray_ints[2000 * 4 + 1]
    #~ asd
    maparray = numpy.zeros(shape=(map.width*map.height*map.numsources, ), dtype=numpy.float32)
    mf = cl.mem_flags
    ray_floats_buf = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=ray_floats)
    ray_ints_buf = cl.Buffer(ctx, mf.COPY_HOST_PTR, hostbuf=ray_ints)
    maparray_buf = cl.Buffer(ctx, mf.COPY_HOST_PTR, hostbuf=maparray)

    prg = cl.Program(ctx, """
        __kernel void sum(__global const float *ray_floats,
        __global const int *ray_ints, __global float *map)
        {
            int gid = get_global_id(0);
            
            
            float dx = ray_floats[gid * 5 + 0];
            float dy = ray_floats[gid * 5 + 1];
            float x = ray_floats[gid * 5 + 2];
            float y = ray_floats[gid * 5 + 3];
            float power = ray_floats[gid * 5 + 4];
            int dist = ray_ints[gid * 4 + 0];
            int source_id = ray_ints[gid * 4 + 1];
            int generation = ray_ints[gid * 4 + 2];
            int medium = ray_ints[gid * 4 + 3];
            
            float newpower, currpower;
            
            int xi, yi, mapidx;
            int width = 860, height = 300;
            int loop;
            int raylength = 1000;
            float _raylength = raylength;
            for(loop=0; loop < raylength; loop++)
            {
                
                newpower = power * (1.0 - (loop / _raylength));
                if (newpower == 0)
                    break;
                
                x = x + dx;
                y = y + dy;
                
                if (x <= 0 or x >= width - 1 or y <= 0 or y >= height)
                    break;
                
                xi = x;
                yi = y;
                
                mapidx = source_id * width * height + yi * width + xi;
                currpower = map[mapidx];
                if (newpower > currpower)
                {
                    map[mapidx] = newpower;
                }
            }
        }
        """).build()

    #~ print map.height, map.width
    #~ return
    print 'num rays: %s' % num_rays
    prg.sum(queue, (num_rays, ), None, ray_floats_buf, ray_ints_buf, maparray_buf)
    cl.enqueue_copy(queue, maparray, maparray_buf)
    #~ print numpy.max(maparray)
    
    #~ rs = maparray.reshape(map.width, map.height, map.numsources)
    #~ a = numpy.transpose(rs, axes=(2,0,1))[0]
    
    #~ pilutil.imsave('d:/out.png', a)
    # load data into map
    map.fromArray(maparray)