# -*- coding: utf-8 -*-
"""
Created on Fri Mar 07 09:52:51 2014

@author: ashiotis
"""
from __future__ import absolute_import
from __future__ import print_function

import sys, numpy, time
from . import utilstest
import fabio
import pyopencl as cl
from pylab import *
from six.moves import input
print("#"*50)
pyFAI = sys.modules["pyFAI"]
from pyFAI import splitPixelFullLUT
from pyFAI import ocl_hist_pixelsplit
#from pyFAI import splitBBoxLUT
#from pyFAI import splitBBoxCSR
#logger = utilstest.getLogger("profile")


ai = pyFAI.load("testimages/halfccd.poni")
data = fabio.open("testimages/halfccd.edf").data

workgroup_size = 256
bins = 1000

pos_in = ai.array_from_unit(data.shape, "corner", unit="2th_deg")

pos = pos_in.reshape(pos_in.size/8,4,2)




ref = splitPixelFullLUT.HistoLUT1dFullSplit(pos,bins, unit="2th_deg")




pos_size = pos.size
#size = data.size
size = pos_size/8

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

d_pos       = cl.array.to_device(queue, pos)
d_preresult = cl.array.empty(queue, (4*workgroup_size,), dtype=numpy.float32)
d_minmax    = cl.array.empty(queue, (4,), dtype=numpy.float32)

with open("../openCL/ocl_lut_pixelsplit.cl", "r") as kernelFile:
    kernel_src = kernelFile.read()

compile_options = "-D BINS=%i  -D NIMAGE=%i -D WORKGROUP_SIZE=%i -D EPS=%e" % \
                (bins, size, workgroup_size, numpy.finfo(numpy.float32).eps)
            
print(compile_options)

program = cl.Program(ctx, kernel_src).build(options=compile_options)

program.reduce1(queue, (workgroup_size*workgroup_size,), (workgroup_size,), d_pos.data,  numpy.uint32(pos_size), d_preresult.data)

program.reduce2(queue, (workgroup_size,), (workgroup_size,), d_preresult.data, d_minmax.data)


min0 = pos[:, :, 0].min()
max0 = pos[:, :, 0].max()
min1 = pos[:, :, 1].min()
max1 = pos[:, :, 1].max()
minmax=(min0,max0,min1,max1)

print(minmax)
print(d_minmax)


memset_size = (bins + workgroup_size - 1) & ~(workgroup_size - 1),

#d_outMax  = cl.array.empty(queue, (bins,), dtype=numpy.int32)

#program.memset_out_int(queue, memset_size, (workgroup_size,), d_outMax.data)

global_size = (size + workgroup_size - 1) & ~(workgroup_size - 1),

#program.lut1(queue, global_size, (workgroup_size,), d_pos.data, d_minmax.data, numpy.uint32(size), d_outMax.data)


#outMax_1  = numpy.copy(d_outMax)



outMax = ref.outMax

idx_ptr = numpy.ndarray(bins+1, dtype=numpy.int32)
idx_ptr[0] = 0
idx_ptr[1:] = outMax.cumsum()



d_idx_ptr = cl.array.to_device(queue, idx_ptr)

#d_lutsize = cl.array.empty(queue, (1,), dtype=numpy.int32)

#program.lut2(queue, (1,), (1,), d_outMax.data, d_idx_ptr.data, d_lutsize.data)

#lutsize  = numpy.ndarray(1, dtype=numpy.int32)

#cl.enqueue_copy(queue, lutsize, d_lutsize.data)

#print lutsize

lut_size = int(idx_ptr[-1])

d_indices  = cl.array.empty(queue, (lut_size,), dtype=numpy.int32)
d_data     = cl.array.empty(queue, (lut_size,), dtype=numpy.float32)

#d_check_atomics = cl.Buffer(ctx, mf.READ_WRITE, 4*lut_size)


d_outMax = cl.array.empty(queue, (bins,), dtype=numpy.int32)

program.memset_out_int(queue, memset_size, (workgroup_size,), d_outMax.data)

d_outData  = cl.array.empty(queue, (bins,), dtype=numpy.float32)
d_outCount = cl.array.empty(queue, (bins,), dtype=numpy.float32)
d_outMerge = cl.array.empty(queue, (bins,), dtype=numpy.float32)

program.lut3(queue, global_size, (workgroup_size,), d_pos.data, d_minmax.data, numpy.uint32(size), d_outMax.data, d_idx_ptr.data, d_indices.data, d_data.data)


outMax_2  = numpy.copy(d_outMax)

#check_atomics = numpy.ndarray(lut_size, dtype=numpy.int32)

#cl.enqueue_copy(queue, check_atomics, d_check_atomics)


program.memset_out(queue, memset_size, (workgroup_size,), d_outData.data, d_outCount.data, d_outMerge.data)




d_image = cl.array.to_device(queue, data)
d_image_float = cl.array.empty(queue, (size,), dtype=numpy.float32)

#program.s32_to_float(queue, global_size, (workgroup_size,), d_image.data, d_image_float)  # Pilatus1M
program.u16_to_float(queue, global_size, (workgroup_size,), d_image.data, d_image_float.data)  # halfccd

program.csr_integrate(queue, (bins*workgroup_size,),(workgroup_size,), d_image_float.data, d_data.data, d_indices.data, d_idx_ptr.data, d_outData.data, d_outCount.data, d_outMerge.data)


#outData  = numpy.ndarray(bins, dtype=numpy.float32)
#outCount = numpy.ndarray(bins, dtype=numpy.float32)
outMerge = numpy.ndarray(bins, dtype=numpy.float32)


#cl.enqueue_copy(queue,outData, d_outData)
#cl.enqueue_copy(queue,outCount, d_outCount)
cl.enqueue_copy(queue,outMerge, d_outMerge.data)

#program.integrate2(queue, (1024,), (workgroup_size,), d_outData, d_outCount, d_outMerge)

#cl.enqueue_copy(queue,outData, d_outData)
#cl.enqueue_copy(queue,outCount, d_outCount)
#cl.enqueue_copy(queue,outMerge, d_outMerge)



ref = ai.integrate1d(data,bins,unit="2th_deg", correctSolidAngle=False, method="splitpixelfull")



#assert(numpy.allclose(ref,outMerge))

plot(ref[0],outMerge, label="ocl_lut_merge")
##plot(ref[0],outData, label="ocl_lut_data")
##plot(ref[0],outCount, label="ocl_lut_count")
plot(*ref, label="ref_merge")
##plot(ref[0], ref[2], label="ref_data")
##plot(ref[0], ref[3], label="ref_count")
###plot(abs(ref-outMerge)/outMerge, label="ocl_csr_fullsplit")
legend()
show()
input()

