#version 460

layout(local_size_x = 16, local_size_y = 8, local_size_z = 1) in;

layout(binding = 0, set = 0) buffer storageBuffer
{
  float imageData[];
};

void main() {
  // The resolution of the buffer, which in this case is a hardcoded vector
  // of 2 unsigned integers:
  const uvec2 resolution = uvec2(800, 600);

  // Get the coordinates of the pixel for this invocation:
  // .-------.-> x
  // |       |
  // |       |
  // '-------'
  // v
  // y
  const uvec2 pixel = gl_GlobalInvocationID.xy;

  // If the pixel is outside of the image, don't do anything:
  if((pixel.x >= resolution.x) || (pixel.y >= resolution.y)) {
    return;
  }
  // Create a vector of 3 floats with a different color per pixel.
  const vec3 pixelColor = vec3(float(pixel.x) / resolution.x,  // Red
                               float(pixel.y) / resolution.y,  // Green
                               0.0);                           // Blue
  // Get the index of this invocation in the buffer:
  uint linearIndex = resolution.x * pixel.y + pixel.x;
  // Write the color to the buffer.
  imageData[3 * linearIndex + 0] = pixelColor.r;
  imageData[3 * linearIndex + 1] = pixelColor.g;
  imageData[3 * linearIndex + 2] = pixelColor.b;
}

#include <nvvk/descriptorsets_vk.hpp>  // For nvvk::DescriptorSetContainer

// Add the required device extensions for Debug Printf. If this is confusing,
// don't worry; we'll remove this in the next chapter.
deviceInfo.addDeviceExtension(VK_KHR_SHADER_NON_SEMANTIC_INFO_EXTENSION_NAME);
VkValidationFeaturesEXT      validationInfo  = nvvk::make<VkValidationFeaturesEXT>();
VkValidationFeatureEnableEXT validationFeatureToEnable 
    = VK_VALIDATION_FEATURE_ENABLE_DEBUG_PRINTF_EXT;
validationInfo.enabledValidationFeatureCount  = 1;
validationInfo.pEnabledValidationFeatures     = &validationFeatureToEnable;
deviceInfo.instanceCreateInfoExt              = &validationInfo;
#ifdef _WIN32
  _putenv_s("DEBUG_PRINTF_TO_STDOUT", "1");
#else   // If not _WIN32
  putenv("DEBUG_PRINTF_TO_STDOUT=1");
#endif  // _WIN32

// Here's the list of bindings for the descriptor set layout, from raytrace.comp.glsl:
// 0 - a storage buffer (the buffer `buffer`)
// That's it for now!
nvvk::DescriptorSetContainer descriptorSetContainer(context);
descriptorSetContainer.addBinding(
    0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT);
    
    
// Create a layout from the list of bindings
descriptorSetContainer.initLayout();

// Create a descriptor pool from the list of 
// bindings with space for 1 set, and allocate that set
descriptorSetContainer.initPool(1);
// Create a simple pipeline layout from the descriptor set layout:
descriptorSetContainer.initPipeLayout();

// Write a single descriptor in the descriptor set.
VkDescriptorBufferInfo descriptorBufferInfo{};
descriptorBufferInfo.buffer = buffer.buffer;    // The VkBuffer object
descriptorBufferInfo.range           
    = bufferSizeBytes;  // The length of memory to bind; offset is 0.
VkWriteDescriptorSet writeDescriptor 
    = descriptorSetContainer.makeWrite(
        0 /*set index*/, 0 /*binding*/, &descriptorBufferInfo);
        
vkUpdateDescriptorSets(context, // The context
    1, &writeDescriptor,  // An array of VkWriteDescriptorSet objects
    0, nullptr);   // An array of VkCopyDescriptorSet objects (unused)

// For the moment, create an empty pipeline layout. You can ignore this code
// for now; we'll replace it in the next chapter.
VkPipelineLayoutCreateInfo pipelineLayoutCreateInfo 
    = nvvk::make<VkPipelineLayoutCreateInfo>();
pipelineLayoutCreateInfo.setLayoutCount             = 0;
pipelineLayoutCreateInfo.pushConstantRangeCount     = 0;
VkPipelineLayout pipelineLayout;
NVVK_CHECK(vkCreatePipelineLayout(context, 
    &pipelineLayoutCreateInfo, VK_NULL_HANDLE, &pipelineLayout));

pipelineCreateInfo.layout = descriptorSetContainer.getPipeLayout();

void vkCmdBindDescriptorSets(
    // The command buffer to record this command to.
    // Subsequent commands in it will use these bindings.
    VkCommandBuffer        commandBuffer,      
	// Where to bind the descriptors (in this case, the compute pipeline bind point).
    // Some commands use bindings from different bind points (e.g. compute dispatches
    // use the compute pipeline bind point, draw calls use the graphics bind point).
    VkPipelineBindPoint    pipelineBindPoint,  
    // The pipeline's pipeline layout (like a function signature, excluding push constant ranges)
    VkPipelineLayout       layout,             
    // (Set to 0) Number to add to the descriptor set index
    uint32_t               firstSet,           
    // Number of descriptor sets in the `pDescriptorSets` array
    uint32_t               descriptorSetCount, 
    // A pointer to an array of descriptor sets
    const VkDescriptorSet* pDescriptorSets,    
    // (Unused) Number of dynamic offsets in the next array.
    uint32_t               dynamicOffsetCount, 
    // (Unused) Pointer to an array of dynamic offsets.
    const uint32_t*        pDynamicOffsets
);   

// Bind the descriptor set
VkDescriptorSet descriptorSet = descriptorSetContainer.getSet(0);
vkCmdBindDescriptorSets(cmdBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, 
    descriptorSetContainer.getPipeLayout(), 0, 1,
    &descriptorSet, 0, nullptr);

descriptorSetContainer.deinit();

#extension GL_EXT_scalar_block_layout : require
...

// The scalar layout qualifier here means to align types according to the alignment
// of their scalar components, instead of e.g. padding them to std140 rules.
layout(binding = 0, set = 0, scalar) buffer storageBuffer
{
  vec3 imageData[];
};
...

imageData[linearIndex] = pixelColor;