/*=========================================================================

   Library: iMSTK

   Copyright (c) Kitware, Inc. & Center for Modeling, Simulation,
   & Imaging in Medicine, Rensselaer Polytechnic Institute.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0.txt

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

=========================================================================*/

#include "imstkVulkanPostProcess.h"

#include "imstkVulkanRenderer.h"

namespace imstk
{
VulkanPostProcess::VulkanPostProcess(VulkanRenderer * renderer, unsigned int level, bool lastPass)
{
    m_lastPass = lastPass;
    m_downsampleLevels = level;
    this->createFramebuffer(renderer, level, lastPass);
}

void
VulkanPostProcess::initialize(VulkanRenderer * renderer, std::string fragmentShaderPath)
{
    this->initializeFramebuffer(renderer);
    this->createDescriptorSetLayouts(renderer);
    this->createFullscreenQuad(renderer);
    this->createPipeline(renderer, fragmentShaderPath);
    this->createDescriptors(renderer);
}

void
VulkanPostProcess::createPipeline(VulkanRenderer * renderer, std::string fragmentSource)
{
    // The vertex shader should be the same for every postprocess
    std::string vertexShaderPath = "./Shaders/VulkanShaders/PostProcessing/postprocess_vert.spv";
    std::ifstream vertexShaderDataStream(vertexShaderPath, std::ios_base::binary);
    char vertexShaderData[16000];
    vertexShaderDataStream.read(vertexShaderData, 15999);
    vertexShaderData[(int)vertexShaderDataStream.gcount()] = '\0';

    VkShaderModuleCreateInfo vertexShaderInfo;
    vertexShaderInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    vertexShaderInfo.pNext = nullptr;
    vertexShaderInfo.flags = 0;
    vertexShaderInfo.codeSize = vertexShaderDataStream.gcount();
    vertexShaderInfo.pCode = (uint32_t *)vertexShaderData;
    if (vkCreateShaderModule(renderer->m_renderDevice, &vertexShaderInfo, nullptr, &m_pipelineComponents.vertexShader) != VK_SUCCESS)
    {
        LOG(FATAL) << "Unable to build vertex shader : " << vertexShaderPath;
    }

    std::string fragmentShaderPath = fragmentSource;
    std::ifstream fragmentShaderDataStream(fragmentShaderPath, std::ios_base::binary);
    char fragmentShaderData[16000];
    fragmentShaderDataStream.read(fragmentShaderData, 15999);
    fragmentShaderData[(int)fragmentShaderDataStream.gcount()] = '\0';

    VkShaderModuleCreateInfo fragmentShaderInfo;
    fragmentShaderInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    fragmentShaderInfo.pNext = nullptr;
    fragmentShaderInfo.flags = 0;
    fragmentShaderInfo.codeSize = fragmentShaderDataStream.gcount();
    fragmentShaderInfo.pCode = (uint32_t *)fragmentShaderData;
    if (vkCreateShaderModule(renderer->m_renderDevice, &fragmentShaderInfo, nullptr, &m_pipelineComponents.fragmentShader) != VK_SUCCESS)
    {
        LOG(FATAL) << "Unable to build fragment shader: " << fragmentShaderPath;
    }

    m_pipelineComponents.shaderInfo.resize(2);

    m_pipelineComponents.fragmentSpecializationInfo.mapEntryCount = 0;
    m_pipelineComponents.fragmentSpecializationInfo.pMapEntries = nullptr;
    m_pipelineComponents.fragmentSpecializationInfo.dataSize = 0;
    m_pipelineComponents.fragmentSpecializationInfo.pData = nullptr;

    // Vertex Shader
    m_pipelineComponents.shaderInfo[0].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    m_pipelineComponents.shaderInfo[0].pNext = nullptr;
    m_pipelineComponents.shaderInfo[0].flags = 0;
    m_pipelineComponents.shaderInfo[0].stage = VK_SHADER_STAGE_VERTEX_BIT;
    m_pipelineComponents.shaderInfo[0].module = m_pipelineComponents.vertexShader;
    m_pipelineComponents.shaderInfo[0].pName = "main";
    m_pipelineComponents.shaderInfo[0].pSpecializationInfo = &m_pipelineComponents.fragmentSpecializationInfo;

    // Fragment Shader
    m_pipelineComponents.shaderInfo[1].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    m_pipelineComponents.shaderInfo[1].pNext = nullptr;
    m_pipelineComponents.shaderInfo[1].flags = 0;
    m_pipelineComponents.shaderInfo[1].stage = VK_SHADER_STAGE_FRAGMENT_BIT;
    m_pipelineComponents.shaderInfo[1].module = m_pipelineComponents.fragmentShader;
    m_pipelineComponents.shaderInfo[1].pName = "main";
    m_pipelineComponents.shaderInfo[1].pSpecializationInfo = &m_pipelineComponents.fragmentSpecializationInfo;

    // Vertex Attributes
    m_pipelineComponents.vertexBindingDescription.resize(1);

    m_pipelineComponents.vertexBindingDescription[0].binding = 0;
    m_pipelineComponents.vertexBindingDescription[0].stride = 4 * 5;
    m_pipelineComponents.vertexBindingDescription[0].inputRate = VK_VERTEX_INPUT_RATE_VERTEX;

    // Vertex Attributes
    m_pipelineComponents.vertexAttributeDescription.resize(2);

    m_pipelineComponents.vertexAttributeDescription[0].location = 0;
    m_pipelineComponents.vertexAttributeDescription[0].binding = 0;
    m_pipelineComponents.vertexAttributeDescription[0].format = VK_FORMAT_R32G32B32_SFLOAT;
    m_pipelineComponents.vertexAttributeDescription[0].offset = 0;

    m_pipelineComponents.vertexAttributeDescription[1].location = 1;
    m_pipelineComponents.vertexAttributeDescription[1].binding = 0;
    m_pipelineComponents.vertexAttributeDescription[1].format = VK_FORMAT_R32G32_SFLOAT;
    m_pipelineComponents.vertexAttributeDescription[1].offset = 4 * 3;

    // Pipeline stages
    m_pipelineComponents.vertexInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
    m_pipelineComponents.vertexInfo.pNext = nullptr;
    m_pipelineComponents.vertexInfo.flags = 0;
    m_pipelineComponents.vertexInfo.vertexBindingDescriptionCount = (uint32_t)m_pipelineComponents.vertexBindingDescription.size();
    m_pipelineComponents.vertexInfo.pVertexBindingDescriptions = &m_pipelineComponents.vertexBindingDescription[0];
    m_pipelineComponents.vertexInfo.vertexAttributeDescriptionCount = (uint32_t)m_pipelineComponents.vertexAttributeDescription.size();
    m_pipelineComponents.vertexInfo.pVertexAttributeDescriptions = &m_pipelineComponents.vertexAttributeDescription[0];

    m_pipelineComponents.inputAssemblyInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
    m_pipelineComponents.inputAssemblyInfo.pNext = nullptr;
    m_pipelineComponents.inputAssemblyInfo.flags = 0;
    m_pipelineComponents.inputAssemblyInfo.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
    m_pipelineComponents.inputAssemblyInfo.primitiveRestartEnable = VK_FALSE;

    m_pipelineComponents.tessellationInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_TESSELLATION_STATE_CREATE_INFO;
    m_pipelineComponents.tessellationInfo.pNext = nullptr;
    m_pipelineComponents.tessellationInfo.flags = 0;
    m_pipelineComponents.tessellationInfo.patchControlPoints = 1; // For now, this should be changed in the future

    m_pipelineComponents.viewports.resize(1);

    m_pipelineComponents.viewports[0].x = 0;
    m_pipelineComponents.viewports[0].y = 0;
    m_pipelineComponents.viewports[0].height = m_framebuffer->m_height;
    m_pipelineComponents.viewports[0].width = m_framebuffer->m_width;
    m_pipelineComponents.viewports[0].minDepth = 0.0;
    m_pipelineComponents.viewports[0].maxDepth = 1.0;

    m_pipelineComponents.scissors.resize(1);

    m_pipelineComponents.scissors[0].offset = { 0, 0 };
    m_pipelineComponents.scissors[0].extent = { m_framebuffer->m_width, m_framebuffer->m_height };

    m_pipelineComponents.viewportInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
    m_pipelineComponents.viewportInfo.pNext = nullptr;
    m_pipelineComponents.viewportInfo.flags = 0;
    m_pipelineComponents.viewportInfo.viewportCount = (uint32_t)m_pipelineComponents.viewports.size();
    m_pipelineComponents.viewportInfo.pViewports = &m_pipelineComponents.viewports[0];
    m_pipelineComponents.viewportInfo.scissorCount = (uint32_t)m_pipelineComponents.scissors.size();
    m_pipelineComponents.viewportInfo.pScissors = &m_pipelineComponents.scissors[0];

    m_pipelineComponents.rasterizationInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
    m_pipelineComponents.rasterizationInfo.pNext = nullptr;
    m_pipelineComponents.rasterizationInfo.flags = 0;
    m_pipelineComponents.rasterizationInfo.depthClampEnable = VK_FALSE;
    m_pipelineComponents.rasterizationInfo.rasterizerDiscardEnable = VK_FALSE;
    m_pipelineComponents.rasterizationInfo.polygonMode = VK_POLYGON_MODE_FILL;
    m_pipelineComponents.rasterizationInfo.cullMode = VK_CULL_MODE_NONE;
    m_pipelineComponents.rasterizationInfo.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
    m_pipelineComponents.rasterizationInfo.depthBiasEnable = VK_FALSE;
    m_pipelineComponents.rasterizationInfo.depthBiasConstantFactor = 0.0;
    m_pipelineComponents.rasterizationInfo.depthBiasClamp = VK_FALSE;
    m_pipelineComponents.rasterizationInfo.depthBiasSlopeFactor = 0.0;
    m_pipelineComponents.rasterizationInfo.lineWidth = 1.0;

    m_pipelineComponents.multisampleInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    m_pipelineComponents.multisampleInfo.pNext = nullptr;
    m_pipelineComponents.multisampleInfo.flags = 0;
    m_pipelineComponents.multisampleInfo.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT; // TODO: Enable multisampling
    m_pipelineComponents.multisampleInfo.sampleShadingEnable = VK_FALSE;
    m_pipelineComponents.multisampleInfo.minSampleShading = 0;
    m_pipelineComponents.multisampleInfo.pSampleMask = nullptr;
    m_pipelineComponents.multisampleInfo.alphaToCoverageEnable = VK_FALSE;
    m_pipelineComponents.multisampleInfo.alphaToOneEnable = VK_FALSE;

    VkStencilOpState states[2];
    states[0].failOp = VK_STENCIL_OP_ZERO;
    states[0].passOp = VK_STENCIL_OP_KEEP;
    states[0].depthFailOp = VK_STENCIL_OP_ZERO;
    states[0].compareOp = VK_COMPARE_OP_LESS;
    states[0].compareMask = 0;
    states[0].writeMask = 0;
    states[0].reference = 0;

    states[1].failOp = VK_STENCIL_OP_ZERO;
    states[1].passOp = VK_STENCIL_OP_KEEP;
    states[1].depthFailOp = VK_STENCIL_OP_ZERO;
    states[1].compareOp = VK_COMPARE_OP_LESS;
    states[1].compareMask = 0;
    states[1].writeMask = 0;
    states[1].reference = 0;

    m_pipelineComponents.depthStencilInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
    m_pipelineComponents.depthStencilInfo.pNext = nullptr;
    m_pipelineComponents.depthStencilInfo.flags = 0;
    m_pipelineComponents.depthStencilInfo.depthTestEnable = VK_FALSE;
    m_pipelineComponents.depthStencilInfo.depthWriteEnable = VK_FALSE;
    m_pipelineComponents.depthStencilInfo.depthCompareOp = VK_COMPARE_OP_LESS;
    m_pipelineComponents.depthStencilInfo.depthBoundsTestEnable = VK_FALSE;
    m_pipelineComponents.depthStencilInfo.stencilTestEnable = VK_FALSE;
    m_pipelineComponents.depthStencilInfo.front = states[0];
    m_pipelineComponents.depthStencilInfo.back = states[1];
    m_pipelineComponents.depthStencilInfo.minDepthBounds = VK_FALSE;
    m_pipelineComponents.depthStencilInfo.maxDepthBounds = VK_FALSE;

    m_pipelineComponents.colorBlendAttachments.resize(m_colorAttachments.size());
    for (int i = 0; i < m_pipelineComponents.colorBlendAttachments.size(); i++)
    {
        m_pipelineComponents.colorBlendAttachments[i].blendEnable = VK_TRUE;
        m_pipelineComponents.colorBlendAttachments[i].srcColorBlendFactor = VK_BLEND_FACTOR_ONE;
        m_pipelineComponents.colorBlendAttachments[i].dstColorBlendFactor = VK_BLEND_FACTOR_ZERO;
        m_pipelineComponents.colorBlendAttachments[i].colorBlendOp = VK_BLEND_OP_ADD;
        m_pipelineComponents.colorBlendAttachments[i].srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
        m_pipelineComponents.colorBlendAttachments[i].dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
        m_pipelineComponents.colorBlendAttachments[i].alphaBlendOp = VK_BLEND_OP_ADD;
        m_pipelineComponents.colorBlendAttachments[i].colorWriteMask = VK_COLOR_COMPONENT_R_BIT |
                                                                       VK_COLOR_COMPONENT_G_BIT |
                                                                       VK_COLOR_COMPONENT_B_BIT |
                                                                       VK_COLOR_COMPONENT_A_BIT;
    }

    m_pipelineComponents.colorBlendInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    m_pipelineComponents.colorBlendInfo.pNext = nullptr;
    m_pipelineComponents.colorBlendInfo.flags = 0;
    m_pipelineComponents.colorBlendInfo.logicOpEnable = VK_FALSE;
    m_pipelineComponents.colorBlendInfo.logicOp = VK_LOGIC_OP_SET;
    m_pipelineComponents.colorBlendInfo.attachmentCount = (uint32_t)m_pipelineComponents.colorBlendAttachments.size();
    m_pipelineComponents.colorBlendInfo.pAttachments = &m_pipelineComponents.colorBlendAttachments[0];
    m_pipelineComponents.colorBlendInfo.blendConstants[0] = 1.0;
    m_pipelineComponents.colorBlendInfo.blendConstants[1] = 1.0;
    m_pipelineComponents.colorBlendInfo.blendConstants[2] = 1.0;
    m_pipelineComponents.colorBlendInfo.blendConstants[3] = 1.0;

    VkPushConstantRange pushConstants;
    pushConstants.offset = 0;
    pushConstants.size = 128; // Minimum on all devices
    pushConstants.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

    VkPipelineLayoutCreateInfo layoutInfo;
    layoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    layoutInfo.pNext = nullptr;
    layoutInfo.flags = 0;
    layoutInfo.setLayoutCount = (uint32_t)m_descriptorSetLayouts.size();
    layoutInfo.pSetLayouts = &m_descriptorSetLayouts[0];
    layoutInfo.pushConstantRangeCount = 1;
    layoutInfo.pPushConstantRanges = &pushConstants;

    vkCreatePipelineLayout(renderer->m_renderDevice, &layoutInfo, nullptr, &m_pipelineLayout);

    m_pipelineComponents.dynamicStates.push_back(VK_DYNAMIC_STATE_VIEWPORT);
    m_pipelineComponents.dynamicStates.push_back(VK_DYNAMIC_STATE_SCISSOR);

    m_pipelineComponents.dynamicStateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    m_pipelineComponents.dynamicStateInfo.pNext = nullptr;
    m_pipelineComponents.dynamicStateInfo.flags = 0;
    m_pipelineComponents.dynamicStateInfo.dynamicStateCount = (uint32_t)m_pipelineComponents.dynamicStates.size();
    m_pipelineComponents.dynamicStateInfo.pDynamicStates = &m_pipelineComponents.dynamicStates[0];

    m_graphicsPipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    m_graphicsPipelineInfo.pNext = nullptr;
    m_graphicsPipelineInfo.flags = VK_PIPELINE_CREATE_DISABLE_OPTIMIZATION_BIT;
    m_graphicsPipelineInfo.stageCount = (uint32_t)m_pipelineComponents.shaderInfo.size();
    m_graphicsPipelineInfo.pStages = &m_pipelineComponents.shaderInfo[0];
    m_graphicsPipelineInfo.pVertexInputState = &m_pipelineComponents.vertexInfo;
    m_graphicsPipelineInfo.pInputAssemblyState = &m_pipelineComponents.inputAssemblyInfo;
    m_graphicsPipelineInfo.pTessellationState = &m_pipelineComponents.tessellationInfo;
    m_graphicsPipelineInfo.pViewportState = &m_pipelineComponents.viewportInfo;
    m_graphicsPipelineInfo.pRasterizationState = &m_pipelineComponents.rasterizationInfo;
    m_graphicsPipelineInfo.pMultisampleState = &m_pipelineComponents.multisampleInfo;
    m_graphicsPipelineInfo.pDepthStencilState = &m_pipelineComponents.depthStencilInfo;
    m_graphicsPipelineInfo.pColorBlendState = &m_pipelineComponents.colorBlendInfo;
    m_graphicsPipelineInfo.pDynamicState = &m_pipelineComponents.dynamicStateInfo;
    m_graphicsPipelineInfo.layout = m_pipelineLayout;
    m_graphicsPipelineInfo.renderPass = m_renderPass;
    m_graphicsPipelineInfo.subpass = 0;
    m_graphicsPipelineInfo.basePipelineHandle = VK_NULL_HANDLE;
    m_graphicsPipelineInfo.basePipelineIndex = 0;
}

void
VulkanPostProcess::createFullscreenQuad(VulkanRenderer * renderer)
{
    m_vertexBuffer = std::make_shared<VulkanVertexBuffer>
                     (renderer->m_memoryManager, 4, (unsigned int)(sizeof(float) * 5), 2);

    {
        auto data = (float*)m_vertexBuffer->getVertexMemory();
        data[0] = -1;
        data[1] = -1;
        data[2] = 0;
        data[3] = 0;
        data[4] = 0;

        data[5] = -1;
        data[6] = 1;
        data[7] = 0;
        data[8] = 0;
        data[9] = 1;

        data[10] = 1;
        data[11] = -1;
        data[12] = 0;
        data[13] = 1;
        data[14] = 0;

        data[15] = 1;
        data[16] = 1;
        data[17] = 0;
        data[18] = 1;
        data[19] = 1;
    }

    {
        auto data = (uint32_t*)m_vertexBuffer->getIndexMemory();
        data[0] = 0;
        data[1] = 1;
        data[2] = 2;
        data[3] = 1;
        data[4] = 2;
        data[5] = 3;
    }

    m_vertexBuffer->initializeBuffers(renderer->m_memoryManager);
}


void
VulkanPostProcess::createDescriptors(VulkanRenderer * renderer)
{
    this->createDescriptorPool(renderer);
    this->createDescriptorSets(renderer);
}

void
VulkanPostProcess::createDescriptorSetLayouts(VulkanRenderer * renderer)
{
    m_descriptorSets.resize(1);
    m_descriptorSetLayouts.resize(1);

    std::vector<VkDescriptorSetLayoutBinding> fragmentDescriptorSetLayoutBindings;

    // Diffuse texture
    for (int i = 0; i < m_samplers.size(); i++)
    {
        VkDescriptorSetLayoutBinding fragmentLayoutBinding;
        fragmentLayoutBinding.binding = i;
        fragmentLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        fragmentLayoutBinding.descriptorCount = 1;
        fragmentLayoutBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;
        fragmentLayoutBinding.pImmutableSamplers = nullptr;
        fragmentDescriptorSetLayoutBindings.push_back(fragmentLayoutBinding);
    }

    VkDescriptorSetLayoutCreateInfo descriptorSetLayoutInfo[1];
    descriptorSetLayoutInfo[0].sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    descriptorSetLayoutInfo[0].pNext = nullptr;
    descriptorSetLayoutInfo[0].flags = 0;
    descriptorSetLayoutInfo[0].bindingCount = (uint32_t)fragmentDescriptorSetLayoutBindings.size();
    descriptorSetLayoutInfo[0].pBindings = &fragmentDescriptorSetLayoutBindings[0];

    for (int i = 0; i < m_descriptorSetLayouts.size(); i++)
    {
        vkCreateDescriptorSetLayout(renderer->m_renderDevice, &descriptorSetLayoutInfo[i], nullptr, &m_descriptorSetLayouts[i]);
    }
}

void
VulkanPostProcess::createDescriptorPool(VulkanRenderer * renderer)
{
    std::vector<VkDescriptorPoolSize> descriptorPoolSizes;

    // Fragment shader textures
    VkDescriptorPoolSize poolSize;
    poolSize.type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    poolSize.descriptorCount = (uint32_t)m_samplers.size();
    descriptorPoolSizes.push_back(poolSize);

    VkDescriptorPoolCreateInfo descriptorPoolInfo;
    descriptorPoolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    descriptorPoolInfo.pNext = nullptr;
    descriptorPoolInfo.flags = 0;
    descriptorPoolInfo.maxSets = (uint32_t)m_descriptorSets.size();
    descriptorPoolInfo.poolSizeCount = (uint32_t)descriptorPoolSizes.size();
    descriptorPoolInfo.pPoolSizes = &descriptorPoolSizes[0];

    vkCreateDescriptorPool(renderer->m_renderDevice, &descriptorPoolInfo, nullptr, &m_descriptorPool);
}

void
VulkanPostProcess::createDescriptorSets(VulkanRenderer * renderer)
{
    VkDescriptorSetAllocateInfo descriptorSetAllocationInfo[1];
    descriptorSetAllocationInfo[0].sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    descriptorSetAllocationInfo[0].pNext = nullptr;
    descriptorSetAllocationInfo[0].descriptorPool = m_descriptorPool;
    descriptorSetAllocationInfo[0].descriptorSetCount = (uint32_t)m_descriptorSetLayouts.size();
    descriptorSetAllocationInfo[0].pSetLayouts = &m_descriptorSetLayouts[0];

    VkDeviceSize size = { VK_WHOLE_SIZE };

    std::vector<VkDescriptorImageInfo> fragmentTextureInfo;

    // Textures
    for (int i = 0; i < m_samplers.size(); i++)
    {
        VkDescriptorImageInfo textureInfo;
        textureInfo.sampler = *m_samplers[i];
        textureInfo.imageView = *m_imageViews[i];
        textureInfo.imageLayout = m_layouts[i];
        fragmentTextureInfo.push_back(textureInfo);
    }

    vkAllocateDescriptorSets(renderer->m_renderDevice, descriptorSetAllocationInfo, &m_descriptorSets[0]);

    m_writeDescriptorSets.resize(1);

    m_writeDescriptorSets[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    m_writeDescriptorSets[0].pNext = nullptr;
    m_writeDescriptorSets[0].dstBinding = 0;
    m_writeDescriptorSets[0].dstArrayElement = 0;
    m_writeDescriptorSets[0].dstSet = m_descriptorSets[0];
    m_writeDescriptorSets[0].descriptorCount = (uint32_t)fragmentTextureInfo.size();
    m_writeDescriptorSets[0].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    m_writeDescriptorSets[0].pBufferInfo = nullptr;
    m_writeDescriptorSets[0].pImageInfo = &fragmentTextureInfo[0];
    m_writeDescriptorSets[0].pTexelBufferView = nullptr;

    vkUpdateDescriptorSets(renderer->m_renderDevice, (uint32_t)m_writeDescriptorSets.size(), &m_writeDescriptorSets[0], 0, nullptr);
}

void
VulkanPostProcess::createRenderPass(VulkanRenderer * renderer)
{
    std::vector<VkAttachmentDescription> attachments;
    bool depth = false;

    // Color attachment
    {
        VkAttachmentDescription attachment;
        attachment.flags = 0;
        attachment.format = m_framebuffer->m_colorFormat;
        attachment.samples = VK_SAMPLE_COUNT_1_BIT;
        attachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
        attachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
        attachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
        attachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
        attachment.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        attachment.finalLayout =
            m_lastPass ? VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL : VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        attachments.push_back(attachment);
    }

    if (m_framebuffer->m_depthFormat != VK_FORMAT_UNDEFINED)
    {
        depth = true;

        // Depth attachment
        VkAttachmentDescription attachment;
        attachment.flags = 0;
        attachment.format = m_framebuffer->m_depthFormat;
        attachment.samples = VK_SAMPLE_COUNT_1_BIT;
        attachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
        attachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
        attachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
        attachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
        attachment.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        attachment.finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
        attachments.push_back(attachment);
    }

    // Color attachment
    VkAttachmentReference colorReference;
    colorReference.attachment = 0;
    colorReference.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

    // Depth attachment
    VkAttachmentReference depthReference;
    depthReference.attachment = 1;
    depthReference.layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

    // Normal attachment
    VkAttachmentReference normalReference;
    normalReference.attachment = 2;
    normalReference.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

    // First pass: geometry
    m_colorAttachments.push_back(colorReference);

    // Render subpasses
    VkSubpassDescription subpassInfo[1];
    subpassInfo[0].flags = 0;
    subpassInfo[0].pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
    subpassInfo[0].inputAttachmentCount = 0;
    subpassInfo[0].pInputAttachments = nullptr;
    subpassInfo[0].colorAttachmentCount = (uint32_t)m_colorAttachments.size();
    subpassInfo[0].pColorAttachments = &m_colorAttachments[0];
    subpassInfo[0].pResolveAttachments = nullptr;
    subpassInfo[0].pDepthStencilAttachment = depth ? &depthReference : nullptr;
    subpassInfo[0].preserveAttachmentCount = 0;
    subpassInfo[0].pPreserveAttachments = nullptr;

    VkSubpassDependency dependencies[2];
    dependencies[0].srcSubpass = VK_SUBPASS_EXTERNAL;
    dependencies[0].dstSubpass = 0;
    dependencies[0].srcStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
    dependencies[0].dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
    dependencies[0].srcAccessMask = VK_ACCESS_SHADER_READ_BIT;
    dependencies[0].dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT | VK_ACCESS_COLOR_ATTACHMENT_READ_BIT;
    dependencies[0].dependencyFlags = VK_DEPENDENCY_BY_REGION_BIT;

    dependencies[1].srcSubpass = 0;
    dependencies[1].dstSubpass = VK_SUBPASS_EXTERNAL;
    dependencies[1].srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
    dependencies[1].dstStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
    dependencies[1].srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT | VK_ACCESS_COLOR_ATTACHMENT_READ_BIT;
    dependencies[1].dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
    dependencies[1].dependencyFlags = VK_DEPENDENCY_BY_REGION_BIT;

    VkRenderPassCreateInfo renderPassInfo;
    renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
    renderPassInfo.pNext = nullptr;
    renderPassInfo.flags = 0;
    renderPassInfo.attachmentCount = (uint32_t)attachments.size();
    renderPassInfo.pAttachments = &attachments[0];
    renderPassInfo.subpassCount = 1;
    renderPassInfo.pSubpasses = subpassInfo;
    renderPassInfo.dependencyCount = 2;
    renderPassInfo.pDependencies = dependencies;

    vkCreateRenderPass(renderer->m_renderDevice, &renderPassInfo, nullptr, &m_renderPass);
}

void
VulkanPostProcess::initializeFramebuffer(VulkanRenderer * renderer)
{
    this->createRenderPass(renderer);
    m_framebuffer->initializeFramebuffer(&m_renderPass);
}

void
VulkanPostProcess::createFramebuffer(VulkanRenderer * renderer,
                                     unsigned int level,
                                     bool lastPass)
{
    m_framebuffer = std::make_shared<VulkanFramebuffer>(
        renderer->m_memoryManager,
        renderer->m_width >> level,
        renderer->m_height >> level,
        lastPass);
}

void
VulkanPostProcess::addInputImage(
    VkSampler * sampler,
    VkImageView * imageView,
    VkImageLayout layout)
{
    m_samplers.push_back(sampler);
    m_imageViews.push_back(imageView);
    m_layouts.push_back(layout);
}

void
VulkanPostProcess::clear(VkDevice * device)
{
    vkDestroyShaderModule(*device, m_pipelineComponents.vertexShader, nullptr);
    vkDestroyShaderModule(*device, m_pipelineComponents.fragmentShader, nullptr);
    vkDestroyPipelineLayout(*device, m_pipelineLayout, nullptr);
    vkDestroyPipeline(*device, m_pipeline, nullptr);
    vkDestroyRenderPass(*device, m_renderPass, nullptr);

    for (auto layout : m_descriptorSetLayouts)
    {
        vkDestroyDescriptorSetLayout(*device, layout, nullptr);
    }

    vkDestroyDescriptorPool(*device, m_descriptorPool, nullptr);
}
}