/**
 * This file contains the compute shaders that make up the logic behind our GPGPU-driven particle system.
 * These are mainly split into the two modules Update, which handles individual particle state updating, and 
 * Emit, which is used to insert new particles into the system.
 * The ComputeIndirectUpdateArgs shader program is used to determine how many thread groups we will need to
 * dispatch the Update shader program on so that it will have at least one thread per existing particle to be
 * updated without the need to involve the CPU.
 * In order for a compute shader program to be recognized for compilation it «must» be given one of the 
 * compute-shader specific group / thread identifier semantics as an input argument to the compute shader function.
 * This holds true even when writing single-thread utility functions like the ComputeIndirectUpdateArgs program
 * at the bottom of this source file.
 * The available semantic types are as follows:
 *	uint3	SV_GroupId				-> The currently executing thread group identified in three dimensions
 *	uint	SV_GroupIndex			-> A flattened version of the above unique identifier of a thread group
 *	uint3	SV_GroupThreadId		-> Identifies the running thread within a thread group
 *	uint3	SV_DispatchThreadId		-> The global thread id within the entire dispatch call; a combination of GroupId and GroupThreadId
 * 
 * You can freely manipulate the number of thread groups to be dispatched during runtime from the DBPro code. 
 * However the number of threads per such group is fixed and must be made available to the shader at compilation time
 * through an attribute declaration like those we saw in the tessellation demo.
 * The number of threads are given as a 3D index (2D on DX10 hardware, in which case the Z dimension must be set to 1).
 * Taken together the number of threads per group cannot exceed 1024 in Shader Model 5 (DX11) or 768 in SM 4 (DX10). The
 * number of thread groups mustn't exceed 65535, while the Z dimension cannot exceed 64 (or 1 for DX10). See the documentation for
 * the DX11 DISPATCH COMPUTE SHADER function for further details about this.
 **/

// Include common declarations
#include "Declarations.hlsl"

// This constant buffer provides per-frame data used by the Emit and Update programs alike
cbuffer GlobalFrameData : register(b0) {
	uint		Time;
	float		DeltaTime;
	float4x4	WorldView;	// This data needs to be available to the Emit / Update compute shader programs
};

// This constant buffer denotes how many particles exist in the Input buffer so that our Update program won't consume past the end of the input buffer
cbuffer UpdateFrameData : register(b1) {
	uint	ParticleCount;
};

// This constant buffer denotes how many new particles should be inserted to the Output buffer by our Emit program
cbuffer EmitFrameData : register(b1) {
	uint	EmitCount;
	uint	RandomSeed;
};

// This is stored in a separate buffer because unlike the EmitFrameData, these constants will not be continously updated during the runtime of our application.
// This allows us to make optimizations so that the values will be accessed faster by the GPU at the cost of making it slower to update the cbuffer from the CPU.
cbuffer EmissionSettings : register(b2) {
	// One of these four emission positions will be picked at random when emitting a new particle
	float3	EmissionPosition[4];
};

// This too is stored in its separate buffer because other shader programs re-use the ParticleCount stored in the UpdateFrameData, while they have no use
// for this particular data that is only used by the Update shader program.
cbuffer ForceFrameData : register(b2) {
	float3	AttractorPos;
	float	AttractionForce;
};

// The Output buffer receives the updated particle data from the Input buffer as well as any newly emitted particles. 
// The Output and Input buffers are then swapped on a per-frame basis.
AppendStructuredBuffer<Particle> Output : register(u0);
// The Input buffer contains the current state of all existing particles prior to this frame's particle state update.
// It is simply the final Output buffer from the previous frame.
ConsumeStructuredBuffer<Particle> Input : register(u1);

// This buffer contains 8192 (by default, since its actual size is queried from the Rand function below the element count can be freely altered from the DBPro 
// code) random float values between 0 and 1 and is used to do pseudo-random lookups based on the current thread id.
Buffer<float> RandomBuf : register(t0);

/**
 * This is a helper function to extract a pseudo-random number in the given range from the RandomBuf.
 * A new seed value will have to be provided at each invocation since we cannot maintain its internal state between threads
 * (actually we could by using another read/write buffer but that is a bit overkill for this application).
 * Take note that the «seed» argument is marked with the "inout" keyword; this means that the function can update the actual variable provided as this parameter.
 * It is used here to ensure that the seed value is incremented so that the next Rand invocation (from the same thread) won't yield the same "random" number again.
 **/
 float Rand(inout uint seed, float fMin, float fMax) {
	// Retrieve the size of the random buffer so that we can wrap the seed value around to ensure we won't go out of its bounds
	uint len;
	RandomBuf.GetDimensions(len);
	// Wrap around the seed value based on the above buffer size so that we create an index to dereference for it
	uint id = seed++ % len;
	// The RandomBuf contains numbers in the range 0..1, randomly scattered by position (see the DBPro source for how this 
	// buffer is filled out). As such we expand the read value to transform it into the desired range.
	return fMin + (RandomBuf[id] * (fMax - fMin));
 }


/************************************************************************************************
 * COMPUTE SHADERS																				*
 ************************************************************************************************/

/**
 * Our Update program is responsible for updating the state of all existing particles.
 * It can also remove particles by not appending them to the Output buffer. In theory new particles can also be inserted from
 * here but we don't do that in this particular implementation to keep it simple.
 **/
[numthreads(64, 1, 1)]
void Update(uint3 threadId : SV_DispatchThreadId) {
	// The gravity force is always applied to all particles
	const float3 Gravity = float3(0, -2.5, 0);
			
	// Because we use a thread count of 64 threads per group and we can only modify how many such thread groups to dispatch during runtime
	// we need to perform a check like this to ensure we won't try to consume a non-existing Particle from the input buffer.
	if(threadId.x < ParticleCount) {
		Particle p = Input.Consume();
		if(p.timeout > Time) {
			// Compute attraction force vector for adding to the particle's current velocity
			float3 vec = AttractorPos - p.pos;
			// Ensure attraction force doesn't exceed a certain threshold
			float3 force = min(length(vec / 16.0), 25.0) * normalize(vec) * AttractionForce;
			
			// Update velocity / position and update the new view depth value of this particular particle
			p.velocity += DeltaTime * (Gravity + force);
			p.pos += DeltaTime * p.velocity;
			float4 viewPos = mul(float4(p.pos, 1), WorldView);
			viewPos /= viewPos.w;
			p.depth = viewPos.z;

			// Finally append the updated particle to the current particle state buffer («Output»)
			if(p.depth > 0.0)
			Output.Append(p);
		}
	}
}

/**
 * Our Emit program inserts new particles and is called after the Update program when the CPU decides a number of new particles are to be added.
 * Each Emit thread adds at most 1 new particle.
 **/
[numthreads(64, 1, 1)]
void Emit(uint3 threadId : SV_DispatchThreadId) {
	// As with the update program we will have AT LEAST enough threads to match the number of particles to emit, but we may have more threads than that.
	// As such we need to ensure we don't create and append more particles than we're supposed to in case we have more threads than required running.
	if(threadId.x < EmitCount) {
		// We need a writable (l-value) variable for use with the inout argument of the Rand() function
		uint seed	= RandomSeed + threadId.x;
		
		Particle p;
		// We want an equal chance for indices 0 .. 3, however not 4
		p.pos		= EmissionPosition[(int)Rand(seed, 0, 3.999999)] + float3(Rand(seed, -3, 3), 0, Rand(seed, -3, 3));
		p.velocity	= float3(Rand(seed, -10.0, 10.0), Rand(seed, 19.0, 22.0), Rand(seed, -10.0, 10.0));
		p.depth		= mul(float4(p.pos, 1), WorldView).z;
		p.size		= Rand(seed, 2.0, 4.0);
		p.colour	= float3(Rand(seed, 0.5, 1), Rand(seed, 0.5, 1), Rand(seed, 0.5, 1));
		p.timeout	= Time + (uint)Rand(seed, 15000, 20000);	// Make this particle exist for anywhere between 15 and 20 seconds
		Output.Append(p);
	}
}

/** 
 * This small single-thread program is used to set up an indirect buffer used to dispatch the Update program on the appropriate number of thread groups.
 * Since we only use the X dimension for simplicity and there are 64 threads per group, this is a very simple calculation. However to avoid the overhead
 * of reading back the current particle count to the CPU this approach is used since this data can be made readily available to the GPU.
 * Indirect buffers are accessed as (raw) byte address buffers by shaders.
 **/
RWByteAddressBuffer IndirectUpdateArgs : register(u0);

[numthreads(1, 1, 1)]
void ComputeIndirectUpdateArgs(uint3 threadId : SV_DispatchThreadId) {
	IndirectUpdateArgs.Store(0, (int)(ParticleCount / 64) + 1);
	// The below will never change and as such really doesn't need to be reset on each invocation of this shader program, but just to be unambiguous they are.
	IndirectUpdateArgs.Store(4, 1);
	IndirectUpdateArgs.Store(8, 1);
}
