#--------------------------------------------#
# pyclaes_ecb.py
#--------------------------------------------#

import numpy
import pyopencl as cl

class claes_ecb(object):

    def __init__(self):
        ''
        self._initOpenCL();


    # TODO: raise proper exceptions
    def _initOpenCL(self):
        ''
        # Get available platforms and choose the first one
        # TODO: select among available platforms?
        platforms = cl.get_platforms();
        if len(platforms) == 0:
            print("Failed to find any OpenCL platform")
            return None

        # Get avaible GPU devices and choose the first one
        # TODO: select among available gpus?
        devices = platforms[0].get_devices(cl.device_type.GPU)
        if len(devices) == 0:
            print("Could not find GPU devices")
            return None
        self.gpuDevice = devices[0]

        # Retrieve useful device informations
        self.wgSize = self.gpuDevice.get_info(cl.device_info.MAX_WORK_GROUP_SIZE)
        self.compUnits = self.gpuDevice.get_info(cl.device_info.MAX_COMPUTE_UNITS)
        self.maxWiSizes = self.gpuDevice.get_info(cl.device_info.MAX_WORK_ITEM_SIZES)
        # TODO: identify Nvidia platform for getting warp size
        #self.warpSize = self.gpuDevice.get_info(cl.device_info.WARP_SIZE_NV)
        self.globMemSize = self.gpuDevice.get_info(cl.device_info.GLOBAL_MEM_SIZE)
        self.locMemSize = self.gpuDevice.get_info(cl.device_info.LOCAL_MEM_SIZE)
        self.constBuffSize = self.gpuDevice.get_info(cl.device_info.MAX_CONSTANT_BUFFER_SIZE)
        self.addrBits = self.gpuDevice.get_info(cl.device_info.ADDRESS_BITS)

        # Create a context using the first GPU device found
        self.context = cl.Context([self.gpuDevice])

        # Finally create the command queue
        self.cmdQueue = cl.CommandQueue(self.context, self.gpuDevice)


    def encrypt(self, key, data):
        ''
        # Get key length and buffer size
        keyLen = len(key)
        buffSize = len(data)

        # Choose optimal parameters
        # TODO: automate choice of:
        # - bulk size;
        # - local work items size;
        # - global work items size

        # Build the kernel
        # TODO: perform only once
        kernelFile = open('aes_ecb.cl', 'r')
        kernelSrc = kernelFile.read()

        program = cl.Program(self.context, kernelSrc)
        program.build(devices=[self.gpuDevice])
        print program

        # Load key and data into device memory
        keyBuffer = cl.Buffer(self.context, cl.mem_flags.READ_ONLY|cl.mem_flags.COPY_HOST_PTR, hostbuf=key)
        inBuffer = cl.Buffer(self.context, cl.mem_flags.READ_WRITE|cl.mem_flags.COPY_HOST_PTR, hostbuf=data)

        # Launch the kernel (awesome :''D)
        # TODO: just testing the correctness of the encryption kernel. Execute only one workitem.
        # After correctness is ensured, use the proper global and local shapes
        program.aes_ecb(self.cmdQueue, (1,1), None, keyBuffer, inBuffer).wait()

        # Read encryption result
        result = numpy.empty_like(inBuffer)
        #result = numpy.zeros(176<<3, dtype=numpy.uint8)
        cl.enqueue_copy(self.cmdQueue, result, inBuffer)
        print result
        return result




if __name__ == '__main__':
    key = numpy.zeros(128, dtype=numpy.uint8)
    data = numpy.zeros(176<<3, dtype=numpy.uint8)


    foo = claes_ecb()
    print('Device info:')
    print('\twgSize: ' + str(foo.wgSize))
    print('\tcompUnits: ' + str(foo.compUnits))
    print('\tmaxWiSizes: ' + str(foo.maxWiSizes))
    print('\tglobMemSize: ' + str(foo.globMemSize))
    print('\tlocMemSize: ' + str(foo.locMemSize))
    print('\tconstBuffSize: ' + str(foo.constBuffSize))
    print('\taddrBits: ' + str(foo.addrBits))

    foo.encrypt(key, data)

    exit(0)