#version 450

layout(local_size_x = 128) in;

struct DrawArraysIndirectArgs
{
  uint count;
  uint instanceCount;
  uint first;
  uint baseInstance;
};

struct ParticleData
{
  vec4 position;
  vec4 orientationSpeedDecay;
  vec4 velocity;
  vec4 acceleration;
  vec4 scaleAndDecay;
  float life;
  uint textureIndex;
};

struct GenericData
{
  //! Delta time
  float dt;

  //! Elapsed time of the application
  float elapsedTime;

  //! Viewport width of the scene (excluding imgui, that means smaller than window)
  uint viewportWidth;

  //! Ditto but for height
  uint viewportHeight;
};

layout(set = 1, binding = 0) uniform CameraData
{
  vec4 position; 
  mat4 vpMat;
  mat4 viewMat;
  mat4 projMat;
} cameraData;


layout (set = 0, binding = 0) uniform GenericDataBuffer
{
  GenericData data;
} genericDataBuffer;

layout (std430, set = 2, binding = 1) coherent restrict readonly buffer ParticlesInputBuffer
{
  ParticleData data[];
} inputParticles;

// output buffer not needed
layout (std430, set = 2, binding = 2) coherent restrict buffer ParticlesOutputBuffer
{
  ParticleData data[];
} outputParticles;

layout (std430, set = 2, binding = 3) coherent restrict buffer ParticlesFreelistBuffer
{
  int freeCount;
  int freeIndices[];

} freelist;

layout (std430, set = 2, binding = 4) coherent restrict buffer IndicesData
{
  uint indices[];
};

layout (std140, set = 2, binding = 5) coherent restrict buffer IndirectDrawArgs
{
  DrawArraysIndirectArgs indirectArgs;
}; 

// push constants
layout(std140, push_constant) uniform EmitterPushConstant
{
  vec4 emitterPosition;
  uint emissionCount;

} emitterPushConstant;

uint pcg_hash(uint seed)
{
  uint state = seed * 747796405u + 2891336453u;
  uint word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
  return (word >> 22u) ^ word;
}

// Used to advance the PCG state.
uint rand_pcg(inout uint rng_state)
{
  uint state = rng_state;
  rng_state = rng_state * 747796405u + 2891336453u;
  uint word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
  return (word >> 22u) ^ word;
}

// Advances the prng state and returns the corresponding random float.
float rand(inout uint state)
{
  uint x = rand_pcg(state);
  state = x;
  return float(x)*uintBitsToFloat(0x2f800004u);
}

void main()
{ 
  uint index = gl_GlobalInvocationID.x;

  ParticleData particle = inputParticles.data[index];

    // Get seed for randomization
  uint pixel_index = uint (particle.position.x + particle.position.y + floatBitsToUint(genericDataBuffer.data.elapsedTime) * (gl_GlobalInvocationID.x + 1));
  uint seed = pcg_hash (pixel_index);

  if (particle.life > 0.0f)
  {
    // update position from velocity
    particle.position += particle.velocity * genericDataBuffer.data.dt;
    particle.velocity += vec4 (rand(seed) * particle.acceleration.x, rand(seed) * particle.acceleration.y, rand(seed) * particle.acceleration.z, 1.0f);
    particle.life -= genericDataBuffer.data.dt;
    
    if (particle.life < 0.0f || particle.scaleAndDecay.x < 0.0f || particle.scaleAndDecay.y < 0.0f)
    {
      particle.life = 0.0f;
      particle.position.x = 9999.0f;
      
      outputParticles.data[index] = particle;
      freelist.freeIndices[atomicAdd(freelist.freeCount, 1)] = int (index);
      return;
    }

    uint drawIndex = atomicAdd (indirectArgs.instanceCount, 1);
    indices[drawIndex] = index;

  }

  outputParticles.data[index] = particle;
}