#pragma kernel CSMain
${VFXGlobalInclude}
${VFXGlobalDeclaration}
${VFXPerPassInclude}
${VFXInclude("Shaders/VFXParticleCommon.template")}

ByteAddressBuffer attributeBuffer;

#if VFX_FEATURE_MOTION_VECTORS
RWByteAddressBuffer elementToVFXBuffer;
#endif

#if VFX_FEATURE_SORT
struct Kvp
{
	float sortKey;
	uint index;
};
#define OutputType Kvp
#else
#define OutputType uint
#endif

#if INDIRECT_BUFFER_COUNT > 0
RWStructuredBuffer<OutputType> outputBuffer0;
#endif
#if INDIRECT_BUFFER_COUNT > 1
RWStructuredBuffer<OutputType> outputBuffer1;
#endif
#if INDIRECT_BUFFER_COUNT > 2
RWStructuredBuffer<OutputType> outputBuffer2;
#endif
#if INDIRECT_BUFFER_COUNT > 3
RWStructuredBuffer<OutputType> outputBuffer3;
#endif
#if INDIRECT_BUFFER_COUNT > 4
#error Too many indirect buffers defined.
#endif

CBUFFER_START(updateParams)
	uint nbMax;
	uint dispatchWidth;
	uint systemSeed;
CBUFFER_END

${VFXGeneratedBlockFunction}

#if VFX_FEATURE_FRUSTUM_CULL || VFX_FEATURE_LOD
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/GeometricTools.hlsl"

bool IsSphereOutsideFrustum(float3 pos, float radius, float4 frustumPlanes[6])
{
    bool outside = false;
    [unroll]
    for (int i = 0; i < 6; ++i)
        outside = outside || DistanceFromPlane(pos, frustumPlanes[i]) < -radius;
    return outside;
}
#endif

[numthreads(NB_THREADS_PER_GROUP,1,1)]
void CSMain(uint3 groupId		: SV_GroupID,
			uint3 groupThreadId	: SV_GroupThreadID)
{
	uint id = groupThreadId.x + groupId.x * NB_THREADS_PER_GROUP + groupId.y * dispatchWidth * NB_THREADS_PER_GROUP;
	if (id < nbMax)
	{
		Attributes attributes = (Attributes)0;
		SourceAttributes sourceAttributes = (SourceAttributes)0;

        uint index = id;
		${VFXLoadAttributes:{alive}}
		if (attributes.alive)
		{
			${VFXLoadAttributes:{(?!(alive))(\b\w)}}
			${VFXProcessBlocks}

            // Recheck alive as blocks can set it to false for manual culling.
            // Test will be stripped if it's not the case anyway.
            if (attributes.alive)
            {
                ${VFXLoadSize}
                float4x4 elementToVFX = GetElementToVFXMatrix(
                    attributes.axisX,
                    attributes.axisY,
                    attributes.axisZ,
                    float3(attributes.angleX,attributes.angleY,attributes.angleZ),
                    float3(attributes.pivotX,attributes.pivotY,attributes.pivotZ),
                    size3,
                    attributes.position);

#if VFX_FEATURE_FRUSTUM_CULL || VFX_FEATURE_LOD
#if VFX_WORLD_SPACE
	            float4x4 elementToWorld = elementToVFX;
                elementToWorld._m03_m13_m23 = GetCameraRelativePositionWS(elementToWorld._m03_m13_m23);
#else
                float4x4 elementToWorld = mul(GetObjectToWorldMatrix(),elementToVFX);          
#endif

	            float xAxisSqrLength = dot(elementToWorld._m00_m10_m20, elementToWorld._m00_m10_m20);
	            float yAxisSqrLength = dot(elementToWorld._m01_m11_m21, elementToWorld._m01_m11_m21);
	            float zAxisSqrLength = dot(elementToWorld._m02_m12_m22, elementToWorld._m02_m12_m22);             
                float radius = 0.5f * sqrt(xAxisSqrLength + yAxisSqrLength + zAxisSqrLength);

                ${VFXLoadParameter:{radiusScale}}
                radius *= radiusScale;

#if VFX_FEATURE_FRUSTUM_CULL
                if (IsSphereOutsideFrustum(elementToWorld._m03_m13_m23, radius, _FrustumPlanes))
                    return;
#endif
#endif          

#if INDIRECT_BUFFER_COUNT > 0
#if VFX_FEATURE_LOD
                uint outputIndex = ~0u;
#if !VFX_FEATURE_FRUSTUM_CULL
                // If particle is out of frustum and frustum culling is disabled, use the lowest LOD
                // This is useful for shadow passes for instance to avoid out of frustum particles to be culled from shadows
                if (IsSphereOutsideFrustum(elementToWorld._m03_m13_m23, radius, _FrustumPlanes))
                    outputIndex = INDIRECT_BUFFER_COUNT - 1;
                else
#endif
                {
                    float viewZ = mul(GetWorldToViewMatrix(), float4(elementToWorld._m03_m13_m23, 1)).z;
                    float4 clip = mul(GetViewToHClipMatrix(), float4(radius, radius, viewZ, 1));
                    float lodValue = max(abs(clip.x),abs(clip.y)) * rcp(max(VFX_EPSILON, clip.w));

                    ${VFXLoadParameter:{lodValues}}
                
                    for (uint i = 0; i < INDIRECT_BUFFER_COUNT; ++i)
                        if (lodValue > lodValues[i])
                        {
                            outputIndex = i;
                            break;
                        }
                }
#elif INDIRECT_BUFFER_COUNT == 1
                uint outputIndex = 0;
#else
                uint outputIndex = attributes.meshIndex;
#endif
                if (outputIndex >= INDIRECT_BUFFER_COUNT)
                    return;


#if VFX_FEATURE_SORT
#if VFX_LOCAL_SPACE
                float3 posRWS = TransformObjectToWorld(attributes.position);
#else
                float3 posRWS = GetCameraRelativePositionWS(attributes.position);
#endif
                float3 camToPos = posRWS - GetCurrentViewPosition();
		
                Kvp output;
                output.sortKey = dot(camToPos,camToPos); // sqr distance to the camera
                output.index = index;
#else
                uint output = index;
#endif

                if (outputIndex == 0)
                {
                    uint indirectIndex = outputBuffer0.IncrementCounter();
			        outputBuffer0[indirectIndex] = output;
                }
#if INDIRECT_BUFFER_COUNT > 1
                else if (outputIndex == 1)
                {
                    uint indirectIndex = outputBuffer1.IncrementCounter();
			        outputBuffer1[indirectIndex] = output;
                }
#if INDIRECT_BUFFER_COUNT > 2
                else if (outputIndex == 2)
                {
                    uint indirectIndex = outputBuffer2.IncrementCounter();
			        outputBuffer2[indirectIndex] = output;
                }
#if INDIRECT_BUFFER_COUNT > 3
                else if (outputIndex == 3)
                {
                    uint indirectIndex = outputBuffer3.IncrementCounter();
			        outputBuffer3[indirectIndex] = output;
                }
#endif
#endif
#endif
#endif

#if VFX_FEATURE_MOTION_VECTORS
#ifdef VFX_FEATURE_MOTION_VECTORS_VERTS
                uint elementToVFXBaseIndex = index * (VFX_FEATURE_MOTION_VECTORS_VERTS * 2 + 1);
#else
                uint elementToVFXBaseIndex = index * 13;
#endif
                elementToVFXBuffer.Store(elementToVFXBaseIndex++ << 2, attributes.alive ? asuint(currentFrameIndex) : 0u);
#ifdef VFX_FEATURE_MOTION_VECTORS_VERTS
                ${VFXMotionVectorVerts}
                UNITY_UNROLL
                for (int itIndexVert = 0; itIndexVert < VFX_FEATURE_MOTION_VECTORS_VERTS - 1; itIndexVert += 2)
                {
                    float4 vertPosA = TransformPositionVFXToNonJitteredClip(verts[itIndexVert]);
                    float4 vertPosB = TransformPositionVFXToNonJitteredClip(verts[itIndexVert + 1]);
                    elementToVFXBuffer.Store4((elementToVFXBaseIndex + itIndexVert * 2) << 2, asuint(float4(vertPosA.xy / vertPosA.w, vertPosB.xy / vertPosB.w)));
                }
                if (VFX_FEATURE_MOTION_VECTORS_VERTS % 2 == 1)
                {
                    int itIndexVert = VFX_FEATURE_MOTION_VECTORS_VERTS - 1;
                    float4 vertPos = TransformPositionVFXToNonJitteredClip(verts[itIndexVert]);
                    elementToVFXBuffer.Store2((elementToVFXBaseIndex + itIndexVert * 2) << 2, asuint(vertPos.xy / vertPos.w));
                }
#else
                UNITY_UNROLL
                for (int itIndexMatrixRow = 0; itIndexMatrixRow < 3; ++itIndexMatrixRow)
                {
                    float4 value = elementToVFX[itIndexMatrixRow] * attributes.alive;
                    elementToVFXBuffer.Store4((elementToVFXBaseIndex + itIndexMatrixRow * 4) << 2, asuint(value));
                }
#endif
#endif
            }
		}		
	}
}