# modified from https://github.com/Erkaman/vulkan_minimal_compute

# Patching vulkan lib:
import collections

collections.Iterable = collections.abc.Iterable
import ctypes
import os

import numpy as np
import vulkan as vk

from radis.gpu.vulkan.pyvkfft_vulkan import prepare_fft

##from pyvkfft_vulkan import prepare_fft


QUERY_POOL_SIZE = 32  # Max number of queries (=timestamps)
DEV_PROPS = vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
HOST_PROPS = (
    vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
)

# TODO: split this up in a purely vulkan class and a radis class
class GPUApplication(object):
    def __init__(self, deviceID=0, path="./", verbose=True):
        # In order to use Vulkan, you must create an instance
        self._instance = None
        self._debugReportCallback = None
        self._physicalDevice = None
        self._device = None
        self._nextDstBinding = 0
        self._timestampLabels = []
        self._bufferObjects = []
        self.verbose = verbose
        # self._deviceID = deviceID
        self._shaderPath = path
        # self._fftApps = {}
        self._fftAppFwd = None
        self._fftAppInv = None
        self._descriptorSetInitialized = False
        self.updateBatchSizeFunctionList = []

        self._computeShaderModules = []
        self._descriptorPool = None
        self._descriptorSetLayout = None
        self._descriptorSet = None
        self._pipelineLayouts = []
        self._pipelines = []

        self._commandPool = None
        self._commandBuffer = None
        self._commandBufferInitialized = False
        self._oneTimeCommandBuffer = None
        self.command_list = []
        self._queryPool = None

        self._enabledLayers = []
        self._enabledExtensions = []

        self._submitInfo = None
        self._queue = None
        self._queueFamilyIndex = -1
        self._memHeapList = []
        self._memTypeDict = {}

        # Initialize vulkan
        self.createInstance()
        self.selectPhysicalDevice(deviceID, verbose=verbose)
        self.createDevice()
        self.createCommandBuffer()

    def free(self):

        self.command_list = []
        # print('Freeing FFT apps... ', end='')

        if self._fftAppFwd is not None:
            self._fftAppFwd.indirectHost = None
            self._fftAppFwd.indirectBuffer = None

        self._indirect_h = None
        self._fftAppFwd = None
        self._fftAppInv = None

        while len(self._pipelines):
            pipeline = self._pipelines.pop()
            vk.vkDestroyPipeline(self._device, pipeline, None)

        while len(self._pipelineLayouts):
            pipelineLayout = self._pipelineLayouts.pop()
            vk.vkDestroyPipelineLayout(self._device, pipelineLayout, None)

        while len(self._computeShaderModules):
            computeShaderModule = self._computeShaderModules.pop()
            vk.vkDestroyShaderModule(self._device, computeShaderModule, None)

        self._descriptorSet = None

        if self._descriptorSetLayout:
            vk.vkDestroyDescriptorSetLayout(
                self._device, self._descriptorSetLayout, None
            )
            self._descriptorSetLayout = None

        if self._descriptorPool:
            vk.vkDestroyDescriptorPool(self._device, self._descriptorPool, None)
            self._descriptorPool = None

        # print(">>> Freeing buffer objects... ")
        while len(self._bufferObjects):
            bufferObject = self._bufferObjects.pop()
            # print(">>>  --" + bufferObject.name)
            bufferObject.free()
        # print(">>> Done!")

        if self._queryPool:
            vk.vkDestroyQueryPool(self._device, self._queryPool, None)
            self._queryPool = None

        if self._fence:
            vk.vkDestroyFence(self._device, self._fence, None)
            self._fence = None

        if self._commandPool:
            vk.vkDestroyCommandPool(self._device, self._commandPool, None)
            self._commandPool = None

        if self._device:
            vk.vkDestroyDevice(self._device, None)
            self._device = None

        if self._instance:
            vk.vkDestroyInstance(self._instance, None)
            self._instance = None

        # print(">>> Freeing Vulkan data...")

    def __setattr__(self, name, val):
        if isinstance(val, GPUBuffer):
            val.app = self
            val.name = name
            val.init_buffer()
            # val._delayedSetData()
        self.__dict__[name] = val

    def cmdScheduleShader(
        self,
        shader_fname=None,
        global_workgroup=(1, 1, 1),
        local_workgroup=(1, 1, 1),
        sync=True,
    ):
        def func(
            self,
            shader_fname=None,
            global_workgroup=(1, 1, 1),
            local_workgroup=(1, 1, 1),
            sync=True,
        ):

            shader_fpath = os.path.join(self._shaderPath, shader_fname)

            # Compute Pipeline:
            pipeline, pipelineLayout, computeShaderModule = self.createComputePipeline(
                shader_fpath, local_workgroup, self._descriptorSetLayout
            )
            self.bindAndDispatch(
                global_workgroup, self._descriptorSet, pipeline, pipelineLayout
            )
            if sync:
                self.sync()

            self._computeShaderModules.append(computeShaderModule)

            self._pipelineLayouts.append(pipelineLayout)
            self._pipelines.append(pipeline)

        return GPUCommand(
            func,
            [
                self,
            ],
            {
                "shader_fname": shader_fname,
                "global_workgroup": global_workgroup,
                "local_workgroup": local_workgroup,
                "sync": sync,
            },
        )

    def cmdScheduleShaderIndirect(
        self,
        shader_fname=None,
        indirect_buffer=None,
        indirect_offset=0,
        local_workgroup=(1, 1, 1),
        sync=True,
    ):
        def func(
            self,
            shader_fname=None,
            indirect_buffer=None,
            indirect_offset=0,
            local_workgroup=(1, 1, 1),
            sync=True,
        ):

            shader_fpath = os.path.join(self._shaderPath, shader_fname)

            # Compute Pipeline:
            pipeline, pipelineLayout, computeShaderModule = self.createComputePipeline(
                shader_fpath, local_workgroup, self._descriptorSetLayout
            )
            self.bindAndDispatchIndirect(
                indirect_buffer,
                indirect_offset,
                self._descriptorSet,
                pipeline,
                pipelineLayout,
            )
            if sync:
                self.sync()

            self._computeShaderModules.append(computeShaderModule)

            self._pipelineLayouts.append(pipelineLayout)
            self._pipelines.append(pipeline)

        return GPUCommand(
            func,
            [
                self,
            ],
            {
                "shader_fname": shader_fname,
                "indirect_buffer": indirect_buffer,
                "indirect_offset": indirect_offset,
                "local_workgroup": local_workgroup,
                "sync": sync,
            },
        )

    def cmdFFT(self, buf, name=""):
        def func(self, buf, name=""):

            if self._fftAppFwd is None:
                self._fftAppFwd = prepare_fft(
                    buf, name=name, compute_app=self, indirectOffset=0, exclusivePlan=1
                )
            self._fftAppFwd.fft(self._commandBuffer, buf._buffer)

        return GPUCommand(func, [self, buf], {"name": name})

    def cmdIFFT(self, buf, name=""):
        def func(self, buf, name=""):

            if self._fftAppInv is None:
                self._fftAppInv = prepare_fft(
                    buf,
                    name=name,
                    compute_app=self,
                    indirectOffset=16 * 4,
                    exclusivePlan=-1,
                )
            self._fftAppInv.ifft(self._commandBuffer, buf._buffer)

        return GPUCommand(func, [self, buf], {"name": name})

    def run(self):
        if self._commandBuffer is None:
            # print('writing command buffer')
            self.writeCommandBuffer()
        # print('cmdbuf:',self._commandBuffer)
        self.runCommandBuffer()

    @staticmethod
    def debugReportCallbackFn(*args):
        print("Debug Report: {} {}".format(args[5], args[6]))
        return 0

    def createInstance(self):
        # Next, we actually create the instance.

        # Contains application info. This is actually not that important.
        # The only real important field is apiVersion.
        applicationInfo = vk.VkApplicationInfo(
            sType=vk.VK_STRUCTURE_TYPE_APPLICATION_INFO,
            pApplicationName="RADIS GPU app",
            applicationVersion=0,
            pEngineName="RADIS Vulkan engine",
            engineVersion=0,
            apiVersion=vk.VK_API_VERSION_1_0,
        )

        createInfo = vk.VkInstanceCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
            flags=0,
            pApplicationInfo=applicationInfo,
            # Give our desired layers and extensions to vulkan.
            enabledLayerCount=len(self._enabledLayers),
            ppEnabledLayerNames=self._enabledLayers,
            enabledExtensionCount=len(self._enabledExtensions),
            ppEnabledExtensionNames=self._enabledExtensions,
        )

        # Actually create the instance.
        # Having created the instance, we can actually start using vulkan.
        self._instance = vk.vkCreateInstance(createInfo, None)

    def selectPhysicalDevice(self, deviceID, verbose=True):
        devices = vk.vkEnumeratePhysicalDevices(self._instance)

        if isinstance(deviceID, str):
            test_str = deviceID.upper()
            for deviceID, device in enumerate(devices):
                props = vk.vkGetPhysicalDeviceProperties(device)
                devname = vk.ffi.string(props.obj.deviceName).decode()
                if test_str in devname.upper():
                    break
            else:
                deviceID = 0
        else:
            if deviceID < 0:
                deviceID += len(devices)
            if not 0 <= deviceID < len(devices):
                deviceID = 0

        self._deviceID = deviceID
        self._physicalDevice = devices[self._deviceID]

        if verbose:
            print("Vulkan version: ", vk.__version__)
            print("Selected device (device_id = {:d}):".format(self._deviceID))
            for i, device in enumerate(devices):
                props = vk.vkGetPhysicalDeviceProperties(device)
                devname = vk.ffi.string(props.obj.deviceName).decode()
                if i == self._deviceID:
                    self._timestampPeriod = props.limits.timestampPeriod

                print(
                    "[{:s}] {:d}: {:s}".format(
                        "X" if i == self._deviceID else " ", i, devname
                    )
                )
            print("")

    # Returns the index of a queue family that supports compute operations.
    def getComputeQueueFamilyIndex(self):
        # Retrieve all queue families.
        queueFamilies = vk.vkGetPhysicalDeviceQueueFamilyProperties(
            self._physicalDevice
        )

        # Now find a family that supports compute.
        for i, props in enumerate(queueFamilies):
            if props.queueCount > 0 and props.queueFlags & vk.VK_QUEUE_COMPUTE_BIT:
                # found a queue with compute. We're done!
                return i

        return -1

    def createDevice(self):
        # We create the logical device in this function.

        self._queueFamilyIndex = self.getComputeQueueFamilyIndex()
        # When creating the device, we also specify what queues it has.
        queueCreateInfo = vk.VkDeviceQueueCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
            queueFamilyIndex=self._queueFamilyIndex,  # find queue family with compute capability.
            queueCount=1,  # create one queue in this family. We don't need more.
            pQueuePriorities=[
                1.0
            ],  # we only have one queue, so this is not that imporant.
        )

        # Now we create the logical device. The logical device allows us to interact with the physical device.
        # Specify any desired device features here. We do not need any for this application, though.
        deviceFeatures = vk.VkPhysicalDeviceFeatures()
        deviceCreateInfo = vk.VkDeviceCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
            enabledLayerCount=len(self._enabledLayers),
            ppEnabledLayerNames=self._enabledLayers,
            pQueueCreateInfos=queueCreateInfo,
            queueCreateInfoCount=1,
            pEnabledFeatures=deviceFeatures,
        )

        self._device = vk.vkCreateDevice(self._physicalDevice, deviceCreateInfo, None)
        self._queue = vk.vkGetDeviceQueue(self._device, self._queueFamilyIndex, 0)

        # We create a fence.
        fenceCreateInfo = vk.VkFenceCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_FENCE_CREATE_INFO, flags=0
        )
        self._fence = vk.vkCreateFence(self._device, fenceCreateInfo, None)
        self._memoryBarrier = vk.VkMemoryBarrier(
            sType=vk.VK_STRUCTURE_TYPE_MEMORY_BARRIER,
            pNext=0,
            srcAccessMask=vk.VK_ACCESS_SHADER_WRITE_BIT,
            dstAccessMask=vk.VK_ACCESS_SHADER_READ_BIT,
        )

        queryPoolCreateInfo = vk.VkQueryPoolCreateInfo(
            vk.VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO,
            None,
            0,
            vk.VK_QUERY_TYPE_TIMESTAMP,
            QUERY_POOL_SIZE,
            0,
        )

        self._queryPool = vk.vkCreateQueryPool(self._device, queryPoolCreateInfo, None)

    def query_memory_properties(self):

        vk_mem_props = {
            "DEVICE_LOCAL": vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
            "HOST_VISIBLE": vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
            "HOST_COHERENT": vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
            "HOST_CACHED": vk.VK_MEMORY_PROPERTY_HOST_CACHED_BIT,
            "LAZILY_ALLOCATED": vk.VK_MEMORY_PROPERTY_LAZILY_ALLOCATED_BIT,
            "DEVICE_COHERENT (AMD)": vk.VK_MEMORY_PROPERTY_DEVICE_COHERENT_BIT_AMD,
            "DEVICE_UNCACHED (AMD)": vk.VK_MEMORY_PROPERTY_DEVICE_UNCACHED_BIT_AMD,
        }

        mem_properties = vk.vkGetPhysicalDeviceMemoryProperties(self._physicalDevice)

        for i in range(mem_properties.memoryHeapCount):
            heap = mem_properties.memoryHeaps[i]
            self._memHeapList.append(heap.size)
            self._memTypeDict[i] = []

        self._memTypeCount = mem_properties.memoryTypeCount
        for i in range(mem_properties.memoryTypeCount):
            mem_type = mem_properties.memoryTypes[i]
            prop_list = []
            for key in vk_mem_props:
                if mem_type.propertyFlags & vk_mem_props[key]:
                    prop_list.append(key)
            self._memTypeDict[mem_type.heapIndex].append([i, prop_list])

    def print_memory_properties(self):
        if not len(self._memHeapList):
            self.query_memory_properties()

        for key in self._memTypeDict.keys():
            print(f"heap {key:d} ({self._memHeapList[key]//(1024*1024):d} MB)")
            for i, (j, l) in enumerate(self._memTypeDict[key]):
                bullet = " └─" if i == len(self._memTypeDict[key]) - 1 else " ├─"
                subbul = "   " if i == len(self._memTypeDict[key]) - 1 else " │ "
                print(bullet + "type:", j)
                for lk in l:
                    bullet = subbul + ("  └─" if lk == l[-1] else "  ├─")
                    print(bullet + lk)
            print("")

    # find memory type with desired properties.
    def findMemoryType2(self, memoryTypeBits, kind):

        memoryTypeBits = 31

        # Look through the heaps, starting from the largest one:
        heap_ids = np.argsort(self._memHeapList)[::-1]
        for heap_id in heap_ids:
            for type, props in self._memTypeDict[heap_id]:
                if not memoryTypeBits & (1 << type):
                    continue

                print(type, props)

        return -1

    # find memory type with desired properties.
    def findLargeMemoryType(self, memoryTypeBits, kind):

        # dev_props = vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
        # host_props = vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
        # dual_props = dev_props | host_props

        mem_properties = vk.vkGetPhysicalDeviceMemoryProperties(self._physicalDevice)

        # First find the largest heap with/without the DEVICE_LOCAL bit:
        heap_dict = {}
        flip = (
            0
            if kind & vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
            else vk.VK_MEMORY_HEAP_DEVICE_LOCAL_BIT
        )
        for i, heap in enumerate(mem_properties.memoryHeaps):
            if (
                heap.flags & vk.VK_MEMORY_HEAP_DEVICE_LOCAL_BIT
            ) ^ flip:  # look for device/host local memory
                heap_dict[heap.size // (1024 * 1024)] = i

        heap_id = heap_dict[
            sorted(heap_dict)[-1]
        ]  # select the biggest heap that is device/host local

        # First check if we can find dual host/device memory:
        for i, type in enumerate(mem_properties.memoryTypes):
            if (type.heapIndex != heap_id) or (not memoryTypeBits & (1 << i)):
                continue

            # If memtype has all requested flags (and possibly others), return
            if (type.propertyFlags & kind) == kind:
                # print(i, type.propertyFlags, kind)
                return i

        return -1

    # find memory type with desired properties.
    def findMemoryType(self, memoryTypeBits, properties):
        memoryProperties = vk.vkGetPhysicalDeviceMemoryProperties(self._physicalDevice)

        for i, mt in enumerate(memoryProperties.memoryTypes):
            if (
                memoryTypeBits & (1 << i)
                and (mt.propertyFlags & properties) == properties
            ):
                return i

        return -1

    def createCommandBuffer(self):
        # We are getting closer to the end. In order to send commands to the device(GPU),
        # we must first record commands into a command buffer.
        # To allocate a command buffer, we must first create a command pool. So let us do that.
        commandPoolCreateInfo = vk.VkCommandPoolCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
            flags=0,
            # the queue family of this command pool. All command buffers allocated from this command pool,
            # must be submitted to queues of this family ONLY.
            queueFamilyIndex=self._queueFamilyIndex,
        )

        self._commandPool = vk.vkCreateCommandPool(
            self._device, commandPoolCreateInfo, None
        )

        # Now allocate a command buffer from the command pool.
        commandBufferAllocateInfo = vk.VkCommandBufferAllocateInfo(
            sType=vk.VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
            commandPool=self._commandPool,
            # if the command buffer is primary, it can be directly submitted to queues.
            # A secondary buffer has to be called from some primary command buffer, and cannot be directly
            # submitted to a queue. To keep things simple, we use a primary command buffer.
            level=vk.VK_COMMAND_BUFFER_LEVEL_PRIMARY,
            commandBufferCount=2,
        )

        self._commandBuffer, self._oneTimeCommandBuffer = vk.vkAllocateCommandBuffers(
            self._device, commandBufferAllocateInfo
        )

        self._commandBufferInitialized = True

    def initDescriptorSet(self):
        # Descriptor set:
        self._descriptorSetLayout = self.createDescriptorSetLayout()
        self._descriptorPool = self.createDescriptorPool()
        self._descriptorSet = self.createDescriptorSet(
            self._descriptorPool, self._descriptorSetLayout
        )

        self._descriptorSetInitialized = True

    def createDescriptorSetLayout(self):
        # Here we specify a descriptor set layout. This allows us to bind our descriptors to
        # resources in the shader.

        # Here we specify a binding of type VK_DESCRIPTOR_TYPE_STORAGE_BUFFER to the binding point
        # 0. This binds to
        #   layout(std140, binding = 0) buffer buf
        # in the compute shader.
        descriptorSetLayoutBindings = [
            vk.VkDescriptorSetLayoutBinding(
                binding=bufferObject._dstBinding,
                descriptorType=bufferObject._descriptorType,
                descriptorCount=1,
                stageFlags=vk.VK_SHADER_STAGE_COMPUTE_BIT,
            )
            for bufferObject in self._bufferObjects
        ]

        descriptorSetLayoutCreateInfo = vk.VkDescriptorSetLayoutCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
            bindingCount=len(descriptorSetLayoutBindings),
            pBindings=descriptorSetLayoutBindings,
        )

        # Create the descriptor set layout.
        return vk.vkCreateDescriptorSetLayout(
            self._device, descriptorSetLayoutCreateInfo, None
        )

    def createDescriptorPool(self):
        descriptorPoolSize = vk.VkDescriptorPoolSize(
            type=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            descriptorCount=len(self._bufferObjects),
        )

        descriptorPoolCreateInfo = vk.VkDescriptorPoolCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
            maxSets=1,  # we only need to allocate one descriptor set from the pool.
            poolSizeCount=1,
            pPoolSizes=descriptorPoolSize,
        )

        # create descriptor pool.
        descriptorPool = vk.vkCreateDescriptorPool(
            self._device, descriptorPoolCreateInfo, None
        )
        return descriptorPool

    def createDescriptorSet(self, descriptorPool, descriptorSetLayout):
        # So we will allocate a descriptor set here.

        # With the pool allocated, we can now allocate the descriptor set.
        descriptorSetAllocateInfo = vk.VkDescriptorSetAllocateInfo(
            sType=vk.VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
            descriptorPool=descriptorPool,
            descriptorSetCount=1,
            pSetLayouts=[descriptorSetLayout],
        )

        # allocate descriptor set.
        descriptorSet = vk.vkAllocateDescriptorSets(
            self._device, descriptorSetAllocateInfo
        )[0]

        return descriptorSet

    def updateDescriptorSet(self, descriptorSet):
        writeDescriptorSets = [
            vk.VkWriteDescriptorSet(
                sType=vk.VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                dstSet=descriptorSet,
                dstBinding=bufferObject._dstBinding,
                descriptorCount=1,
                descriptorType=bufferObject._descriptorType,
                pBufferInfo=vk.VkDescriptorBufferInfo(
                    buffer=bufferObject._buffer,
                    offset=0,
                    range=bufferObject._bufferSize,
                ),
            )
            for bufferObject in self._bufferObjects
            if bufferObject._descriptorType
        ]

        # perform the update of the descriptor set.
        vk.vkUpdateDescriptorSets(
            self._device, len(writeDescriptorSets), writeDescriptorSets, 0, None
        )

    def createComputePipeline(
        self, shaderFileName, localWorkGroup, descriptorSetLayout
    ):
        # We create a compute pipeline here.

        # Create a shader module. A shader module basically just encapsulates some shader code.
        with open(shaderFileName, "rb") as comp:
            code = comp.read()

            createInfo = vk.VkShaderModuleCreateInfo(
                sType=vk.VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
                codeSize=len(code),
                pCode=code,
            )

            computeShaderModule = vk.vkCreateShaderModule(
                self._device, createInfo, None
            )

        # specialize local workgroup size:
        dtype = "int"
        buffer = vk.ffi.new(dtype + "[]", localWorkGroup)
        pData = vk.ffi.cast("void*", buffer)

        dsize = vk.ffi.sizeof(dtype)
        entries = [
            vk.VkSpecializationMapEntry(constantID=i, offset=i * dsize, size=dsize)
            for i in range(len(buffer))
        ]

        specializationInfo = vk.VkSpecializationInfo(
            mapEntryCount=len(entries),
            pMapEntries=entries,
            dataSize=vk.ffi.sizeof(buffer),
            pData=pData,
        )

        # Now let us actually create the compute pipeline.
        # A compute pipeline is very simple compared to a graphics pipeline.
        # It only consists of a single stage with a compute shader.
        # So first we specify the compute shader stage, and it's entry point(main).
        shaderStageCreateInfo = vk.VkPipelineShaderStageCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
            stage=vk.VK_SHADER_STAGE_COMPUTE_BIT,
            module=computeShaderModule,
            pSpecializationInfo=specializationInfo,
            pName="main",
        )

        # The pipeline layout allows the pipeline to access descriptor sets.
        # So we just specify the descriptor set layout we created earlier.
        pipelineLayoutCreateInfo = vk.VkPipelineLayoutCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
            setLayoutCount=1,
            pSetLayouts=[descriptorSetLayout],
        )
        pipelineLayout = vk.vkCreatePipelineLayout(
            self._device, pipelineLayoutCreateInfo, None
        )

        pipelineCreateInfo = vk.VkComputePipelineCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
            stage=shaderStageCreateInfo,
            layout=pipelineLayout,
        )

        # Now, we finally create the compute pipeline.
        pipelines = vk.vkCreateComputePipelines(
            self._device, vk.VK_NULL_HANDLE, 1, pipelineCreateInfo, None
        )
        # print(self._device, pipelines)
        if len(pipelines) == 1:
            pipeline = pipelines[0]
        else:
            raise Exception("Could not create compute pipeline")

        return pipeline, pipelineLayout, computeShaderModule

    def bindAndDispatch(self, globalWorkGroup, descriptorSet, pipeline, pipelineLayout):
        # We need to bind a pipeline, AND a descriptor set before we dispatch.
        # The validation layer will NOT give warnings if you forget these, so be very careful not to forget them.
        vk.vkCmdBindPipeline(
            self._commandBuffer, vk.VK_PIPELINE_BIND_POINT_COMPUTE, pipeline
        )
        vk.vkCmdBindDescriptorSets(
            self._commandBuffer,
            vk.VK_PIPELINE_BIND_POINT_COMPUTE,
            pipelineLayout,
            0,
            1,
            [descriptorSet],
            0,
            None,
        )

        # Calling vkCmdDispatch basically starts the compute pipeline, and executes the compute shader.
        # The number of workgroups is specified in the arguments.
        # If you are already familiar with compute shaders from OpenGL, this should be nothing new to you.
        vk.vkCmdDispatch(self._commandBuffer, *globalWorkGroup)

    # def endCommandBuffer(self):
    # vk.vkEndCommandBuffer(self._commandBuffer)

    def bindAndDispatchIndirect(
        self, indirect_buffer, indirect_offset, descriptorSet, pipeline, pipelineLayout
    ):
        # We need to bind a pipeline, AND a descriptor set before we dispatch.
        # The validation layer will NOT give warnings if you forget these, so be very careful not to forget them.
        vk.vkCmdBindPipeline(
            self._commandBuffer, vk.VK_PIPELINE_BIND_POINT_COMPUTE, pipeline
        )
        vk.vkCmdBindDescriptorSets(
            self._commandBuffer,
            vk.VK_PIPELINE_BIND_POINT_COMPUTE,
            pipelineLayout,
            0,
            1,
            [descriptorSet],
            0,
            None,
        )

        # Calling vkCmdDispatch basically starts the compute pipeline, and executes the compute shader.
        # The number of workgroups is specified in the arguments.
        # If you are already familiar with compute shaders from OpenGL, this should be nothing new to you.
        vk.vkCmdDispatchIndirect(
            self._commandBuffer, indirect_buffer._buffer, indirect_offset
        )

    # def endCommandBuffer(self):
    # vk.vkEndCommandBuffer(self._commandBuffer)

    def createBuffer(
        self,
        bufferSize,
        kind=vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
        descriptorType=vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        usage=vk.VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
        sharingMode=vk.VK_SHARING_MODE_EXCLUSIVE,
    ):
        # create buffer:
        bufferCreateInfo = vk.VkBufferCreateInfo(
            sType=vk.VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
            size=bufferSize,  # buffer size in bytes.
            usage=usage,
            sharingMode=sharingMode,  # buffer is exclusive to a single queue family at a time.
        )
        buffer = vk.vkCreateBuffer(self._device, bufferCreateInfo, None)

        # allocate memory:
        memoryRequirements = vk.vkGetBufferMemoryRequirements(self._device, buffer)
        if (
            usage & vk.VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT
            == vk.VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT
            or usage & vk.VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT
            == vk.VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT
        ):
            # if looking for an indirect buffer, prioritize a host visible device buffer
            index = self.findMemoryType(memoryRequirements.memoryTypeBits, kind)
        else:
            # otherwise, prioritize the largest buffer, possible exclusively device local:
            index = self.findLargeMemoryType(memoryRequirements.memoryTypeBits, kind)

        # print(memoryRequirements.memoryTypeBits)
        # print(' Found type ',index)

        if index < 0:
            vk.vkDestroyBuffer(self._device, buffer, None)
            return None, None, index

        # print('memtype found:', index)

        allocateInfo = vk.VkMemoryAllocateInfo(
            sType=vk.VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
            allocationSize=memoryRequirements.size,
            memoryTypeIndex=index,
        )

        bufferMemory = vk.vkAllocateMemory(self._device, allocateInfo, None)

        # bind buffer:
        vk.vkBindBufferMemory(self._device, buffer, bufferMemory, 0)

        # self._pmappedMemory = vk.vkMapMemory(
        # self._device, bufferMemory, 0, bufferSize, 0
        # )
        # self._hostPtr = int(vk.ffi.cast("unsigned long long", vk.ffi.from_buffer(self._pmappedMemory)))
        # self._bufferObjects.append(self)

        # self._isInitialized = True
        return buffer, bufferMemory, index

    def writeCommandBuffer(self):

        if not self._descriptorSetInitialized:
            # print('Initializing descriptor set... ', end='')
            self.initDescriptorSet()
            # print('Done!')

        # print('Updating descriptor set... ', end='')
        self.updateDescriptorSet(self._descriptorSet)
        # print('Done!')

        if self._commandBuffer is None:
            # print('Creating command buffer... ',end='')
            self.createCommandBuffer()
            # print('Done!')
        else:
            # print('Resetting command buffer... ', end='')
            vk.vkQueueWaitIdle(self._queue)
            vk.vkResetCommandBuffer(self._commandBuffer, 0)
            self._timestampLabels = []
            # print('Done!')

        # Now we shall start recording commands into the newly allocated command buffer.
        beginInfo = vk.VkCommandBufferBeginInfo(
            sType=vk.VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
            flags=0,  # VK_COMMAND_BUFFER_USAGE_SIMULTANEOUS_USE_BIT #VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT
        )
        vk.vkBeginCommandBuffer(self._commandBuffer, beginInfo)
        for obj in self.command_list:
            obj.writeCommand()

        vk.vkEndCommandBuffer(self._commandBuffer)
        # print('Command buffer written!')

    def runCommandBuffer(self):
        # Now we shall finally submit the recorded command buffer to a queue.
        if self._submitInfo is None:
            self._submitInfo = vk.VkSubmitInfo(
                sType=vk.VK_STRUCTURE_TYPE_SUBMIT_INFO,
                commandBufferCount=1,  # submit a single command buffer
                pCommandBuffers=[self._commandBuffer],  # the command buffer to submit.
            )

        # We submit the command buffer on the queue, at the same time giving a fence.
        vk.vkQueueSubmit(self._queue, 1, self._submitInfo, self._fence)

        # The command will not have finished executing until the fence is signalled.
        # So we wait here.
        # We will directly after this read our buffer from the GPU,
        # and we will not be sure that the command has finished executing unless we wait for the fence.
        # Hence, we use a fence here.
        vk.vkWaitForFences(self._device, 1, [self._fence], vk.VK_TRUE, 100000000000)
        vk.vkResetFences(self._device, 1, [self._fence])

    def freeCommandBuffer(self, wait=True):
        if wait:
            vk.vkQueueWaitIdle(self._queue)

        # self.command_list = []

        vk.vkFreeCommandBuffers(
            device=self._device,
            commandPool=self._commandPool,
            commandBufferCount=1,
            pCommandBuffers=[self._commandBuffer],
        )
        # print('Command buffer freed!')

    def oneTimeCommand(self, command, *vargs, **kwargs):

        # wait for queue
        vk.vkQueueWaitIdle(self._queue)
        # reset command buffer
        # vk.vkResetCommandBuffer(command_buffer, 0) # Reset but don't return resources

        # Now we shall start recording commands into the newly allocated command buffer.
        beginInfo = vk.VkCommandBufferBeginInfo(
            sType=vk.VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
            flags=vk.VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
        )
        # vk.vkResetCommandBuffer(self._oneTimeCommandBuffer, 0)
        vk.vkBeginCommandBuffer(self._oneTimeCommandBuffer, beginInfo)
        command(self._oneTimeCommandBuffer, *vargs, **kwargs)
        vk.vkEndCommandBuffer(self._oneTimeCommandBuffer)

        # Now we shall finally submit the recorded command buffer to a queue.
        submitInfo = vk.VkSubmitInfo(
            sType=vk.VK_STRUCTURE_TYPE_SUBMIT_INFO,
            commandBufferCount=1,  # submit a single command buffer
            pCommandBuffers=[
                self._oneTimeCommandBuffer
            ],  # the command buffer to submit.
        )
        # We submit the command buffer on the queue, at the same time giving a fence.
        vk.vkQueueSubmit(self._queue, 1, submitInfo, self._fence)

        # The command will not have finished executing until the fence is signalled.
        # So we wait here.
        # We will directly after this read our buffer from the GPU,
        # and we will not be sure that the command has finished executing unless we wait for the fence.
        # Hence, we use a fence here.
        vk.vkWaitForFences(self._device, 1, [self._fence], vk.VK_TRUE, 100000000000)
        vk.vkResetFences(self._device, 1, [self._fence])

    def cmdClearBuffer(self, buffer_obj):
        def func(self, buffer_obj):
            vk.vkCmdFillBuffer(
                self._commandBuffer, buffer_obj._buffer, 0, buffer_obj._bufferSize, 0
            )

        return GPUCommand(func, [self, buffer_obj], {})

    def sync(self):
        vk.vkCmdPipelineBarrier(
            self._commandBuffer,
            vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            0,
            1,
            self._memoryBarrier,
            0,
            0,
            0,
            0,
        )

    def cmdAddTimestamp(self, label=None):
        def func(self, label=None):
            query = len(self._timestampLabels)
            if label is None:
                label = "timestamp_{:d}".format(query)
            # TODO: deal with duplicates
            self._timestampLabels.append(label)

            if query == 0:
                vk.vkCmdResetQueryPool(
                    self._commandBuffer, self._queryPool, 0, QUERY_POOL_SIZE
                )

            vk.vkCmdWriteTimestamp(
                self._commandBuffer,
                vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
                self._queryPool,
                query,
            )

            if query == QUERY_POOL_SIZE + 1:
                print("Max query pool size reached!!!")
                # TODO: handle exception

        return GPUCommand(
            func,
            [
                self,
            ],
            {"label": label},
        )

    def get_timestamps(self):

        dtype = "unsigned int"
        dsize = vk.ffi.sizeof(dtype)
        length = len(self._timestampLabels)
        queryResult = vk.ffi.new(dtype + "[{:d}]".format(length))
        pData = vk.ffi.cast("void*", queryResult)

        vk.vkGetQueryPoolResults(
            self._device, self._queryPool, 0, length, dsize * length, pData, dsize, 0
        )

        result_arr = np.array([*queryResult]) * self._timestampPeriod * 1e-6  # in ms
        result_arr -= result_arr[0]
        result_dict = dict(
            zip(self._timestampLabels[:-1], result_arr[1:] - result_arr[:-1])
        )
        result_dict["Total"] = result_arr[-1]

        return result_dict

    def setIndirectBuffer(self, buffer, host_type):
        # self._fftAppFwd = None #Reset forward FFT app to prevent an old app referencing an old buffer
        indirect_h = buffer.getHostStructPtr(host_type)
        ctypes.memset(ctypes.byref(indirect_h), 0, ctypes.sizeof(indirect_h))
        self._indirect_h = indirect_h
        return indirect_h

    def setBatchSize(self, N):
        for f in self.updateBatchSizeFunctionList:
            f(N)

    def setFwdFFTWorkGroupSize(self, N):
        self._fftAppFwd.setFFTWorkGroupSize(N)

    def registerBatchSizeUpdateFunction(self, f):
        self.updateBatchSizeFunctionList.append(f)

    def appendCommands(self, cmd):
        if isinstance(cmd, list):
            self.command_list += cmd
        else:
            self.command_list.append(cmd)


class GPUCommand:
    def __init__(self, func, vargs, kwargs):
        self.func = func
        self.vargs = vargs
        self.kwargs = kwargs

    def writeCommand(self):
        self.func(*self.vargs, **self.kwargs)


class GPUBuffer:
    def __init__(
        self,
        bufferSize=0,
        fftSize=None,
        usage="storage",
        binding=None,
        dtype=None,
        app=None,
    ):
        # customization:

        if usage == "indirect":
            self._usage = vk.VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT
            self._descriptorType = None
        elif usage == "uniform":
            self._usage = vk.VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT
            self._descriptorType = vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER
        else:  # usage == 'storage':
            self._usage = vk.VK_BUFFER_USAGE_STORAGE_BUFFER_BIT
            self._descriptorType = vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER

        if fftSize is not None:
            fftSize = np.atleast_1d(fftSize)
            self._batchSize = 1 if len(fftSize) == 1 else fftSize[-2]
            self._fftSize = fftSize[-1]
            self._dtype = np.dtype(np.float32 if dtype is None else dtype)
            self._itemsize = self._dtype.itemsize
            self._bufferSize = (
                self._batchSize * (self._fftSize // 2 + 1) * 2 * self._itemsize
            )
            # self._shape = (self._batchSize, self._fftSize)
        else:
            self._bufferSize = bufferSize

        self._dstBinding = binding

        self._isInitialized = False
        self._stagingBufferInitialized = False

        self.name = ""
        self.app = app
        if self.app is not None:
            self.init_buffer()

    def init_buffer(self):
        self._device = self.app._device

        if self._dstBinding is None:
            self._dstBinding = self.app._nextDstBinding
        self.app._nextDstBinding = self._dstBinding + 1

        self._buffer, self._bufferMemory, index = self.app.createBuffer(
            self._bufferSize,
            kind=(
                vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
                | vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT
                | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
            ),
            descriptorType=self._descriptorType,
            usage=self._usage,
            # sharingMode=vk.VK_SHARING_MODE_EXCLUSIVE,
        )
        self._combined = True

        if index < 0:
            # apparently we don't have memory type that is both host_visible& host_coherent, as well as device_local.
            # this is the case on dedicated GPU's. In this case we only look for device_local.

            self._buffer, self._bufferMemory, index = self.app.createBuffer(
                self._bufferSize,
                kind=vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
                descriptorType=self._descriptorType,
                usage=self._usage,
                # sharingMode=vk.VK_SHARING_MODE_EXCLUSIVE,
            )
            self._combined = False

        if self not in self.app._bufferObjects:
            self.app._bufferObjects.append(self)

        self._isInitialized = True

    def initStagingBuffer(self, chunksize=None):

        if not self._isInitialized:
            print(
                "buffer {self._name} not initialized! First execute self.init_buffer()!"
            )
            return -1

        # self._combined should tell us whether we have a combined host/device buffer

        if self._combined:  # we're dealing with a buffer that is both
            self._stagingBufferSize = self._bufferSize
            self._stagingBuffer = self._buffer
            self._stagingBufferMemory = self._bufferMemory

        else:
            chunksize = self._bufferSize if chunksize is None else chunksize
            self._stagingBufferSize = chunksize

            self._stagingBuffer, self._stagingBufferMemory, _ = self.app.createBuffer(
                self._stagingBufferSize,
                kind=vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT
                | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
                descriptorType=self._descriptorType,
                usage=self._usage,
            )

        self._pmappedMemory = vk.vkMapMemory(
            self._device, self._stagingBufferMemory, 0, self._stagingBufferSize, 0
        )
        self._hostPtr = int(
            vk.ffi.cast("unsigned long long", vk.ffi.from_buffer(self._pmappedMemory))
        )

        self._stagingBufferInitialized = True

    def copyToBuffer(self, arr, device_offset=0, chunksize=None):
        if not self._stagingBufferInitialized:
            self.initStagingBuffer(chunksize=chunksize)

        data_size = arr.nbytes
        dstOffset = device_offset if self._combined else 0
        offset = 0
        while offset < data_size:
            copy_size = min(self._stagingBufferSize, data_size - offset)
            self.fromArray(
                arr, srcOffset=offset, dstOffset=dstOffset, copy_size=copy_size
            )
            self.transferStagingBuffer(
                direction="H2D", dstOffset=offset + device_offset, size=copy_size
            )
            offset += copy_size

        return offset  # returns the amount of bytes copied

    def copyFromBuffer(self, arr, device_offset=0, chunksize=None):
        if not self._stagingBufferInitialized:
            self.initStagingBuffer(chunksize=chunksize)

        data_size = arr.nbytes
        srcOffset = device_offset if self._combined else 0
        offset = 0
        while offset < data_size:
            copy_size = min(self._stagingBufferSize, data_size - offset)
            self.transferStagingBuffer(
                direction="D2H", srcOffset=offset + device_offset, size=copy_size
            )
            self.toArray(
                arr, srcOffset=srcOffset, dstOffset=offset, copy_size=copy_size
            )
            offset += copy_size

        return offset  # returns the amount of bytes copied

    def toArray(self, arr, srcOffset=0, dstOffset=0, copy_size=None):
        if not self._stagingBufferInitialized:
            self.initStagingBuffer()

        if copy_size is None:
            copy_size = arr.nbytes
        ctypes.memmove(
            arr.ctypes.data + dstOffset,
            ctypes.c_void_p(self._hostPtr + srcOffset),
            copy_size,
        )
        return copy_size

    def fromArray(self, arr, srcOffset=0, dstOffset=0, copy_size=None):
        if not self._stagingBufferInitialized:
            self.initStagingBuffer()

        if copy_size is None:
            copy_size = arr.nbytes
        ctypes.memmove(
            ctypes.c_void_p(self._hostPtr + dstOffset),
            arr.ctypes.data + srcOffset,
            copy_size,
        )
        return copy_size

    def transferStagingBuffer(
        self, direction="H2D", srcOffset=0, dstOffset=0, size=None, transfer_now=True
    ):

        if not self._stagingBufferInitialized:
            self.initStagingBuffer()

        if self._combined:
            return

        if direction == "H2D":
            srcBuffer = self._stagingBuffer
            dstBuffer = self._buffer

        elif direction == "D2H":
            srcBuffer = self._buffer
            dstBuffer = self._stagingBuffer

        else:
            return

        if size is None:
            size = self._stagingBufferSize

        copy_region = vk.VkBufferCopy(
            srcOffset=srcOffset, dstOffset=dstOffset, size=size
        )

        if transfer_now:
            self.app.oneTimeCommand(
                vk.vkCmdCopyBuffer,
                srcBuffer=srcBuffer,
                dstBuffer=dstBuffer,
                regionCount=1,
                pRegions=[copy_region],
            )

        else:
            vk.vkCmdCopyBuffer(
                self.app._commandBuffer,
                srcBuffer=srcBuffer,
                dstBuffer=dstBuffer,
                regionCount=1,
                pRegions=[copy_region],
            )

    def cmdTransferStagingBuffer(
        self, direction="H2D", srcOffset=0, dstOffset=0, size=None
    ):
        # TODO: add timestamp , timestamp=False
        return GPUCommand(
            self.transferStagingBuffer,
            [],
            {
                "direction": direction,
                "srcOffset": srcOffset,
                "dstOffset": dstOffset,
                "size": size,
                "transfer_now": False,
            },
        )

    def cmdClearBuffer(self):
        return self.app.cmdClearBuffer(self)

    def setBatchSize(self, batch, grow_only=True, factor=1.5):
        # print('batch:', batch)
        if batch > self._batchSize:
            self._batchSize = int(factor * batch)
            new_size = self._batchSize * (self._fftSize // 2 + 1) * 2 * self._itemsize
            self.resizeBuffer(new_size)

    def resizeBuffer(self, nbytes):
        # print('resizing buffer from',self._bufferSize, 'to', nbytes)
        self.app._fftAppFwd = None
        self.free()
        self._bufferSize = nbytes
        self.init_buffer()

        if self.app._commandBuffer is not None:  # only rewrite if one already exists
            self.app.writeCommandBuffer()

    def getHostStructPtr(self, struct):
        if not self._stagingBufferInitialized:
            self.initStagingBuffer()
        self._structPtr = ctypes.cast(self._hostPtr, ctypes.POINTER(struct))
        return self._structPtr.contents

    def free(self):
        self._isInitialized = False

        if self._stagingBufferInitialized:
            self._stagingBufferInitialized = False

            vk.vkUnmapMemory(self._device, self._stagingBufferMemory)
            self._hostPtr = None
            self._structPtr = None

            if not self._combined:
                vk.vkDestroyBuffer(self._device, self._stagingBuffer, None)
                vk.vkFreeMemory(self._device, self._stagingBufferMemory, None)
            self._stagingBufferMemory = None
            self._stagingBuffer = None
            self._stagingBufferSize = 0

        if self._buffer:
            vk.vkDestroyBuffer(self._device, self._buffer, None)
        if self._bufferMemory:
            vk.vkFreeMemory(self._device, self._bufferMemory, None)

        self._buffer = None
        self._bufferMemory = None
        self._bufferSize = 0
