// Copyright (C) 2009-2021, Panagiotis Christopoulos Charitos and contributors.
// All rights reserved.
// Code licensed under the BSD License.
// http://www.anki3d.org/LICENSE

#pragma anki start comp

ANKI_SPECIALIZATION_CONSTANT_UVEC2(FB_SIZE, 0u);
ANKI_SPECIALIZATION_CONSTANT_UVEC2(TILE_COUNTS, 2u);
ANKI_SPECIALIZATION_CONSTANT_U32(Z_SPLIT_COUNT, 4u);
ANKI_SPECIALIZATION_CONSTANT_U32(TILE_SIZE, 5u);

#define CLUSTERED_SHADING_SET 0
#define CLUSTERED_SHADING_UNIFORMS_BINDING 0
#define CLUSTERED_SHADING_LIGHTS_BINDING 1
#define CLUSTERED_SHADING_CLUSTERS_BINDING 4
#include <AnKi/Shaders/ClusteredShadingCommon.glsl>

const UVec2 WORKGROUP_SIZE = UVec2(8, 8);
layout(local_size_x = WORKGROUP_SIZE.x, local_size_y = WORKGROUP_SIZE.y, local_size_z = 1) in;

layout(set = 0, binding = 5, rgba8) writeonly uniform image2D u_outImg;
layout(set = 0, binding = 6) uniform sampler u_linearAnyClampSampler;
layout(set = 0, binding = 7) uniform texture2D u_depthRt;

void main()
{
	if(skipOutOfBoundsInvocations(WORKGROUP_SIZE, FB_SIZE))
	{
		return;
	}

	// World position
	const Vec2 uv = (Vec2(gl_GlobalInvocationID.xy) + 0.5) / Vec2(FB_SIZE);
	const Vec2 ndc = UV_TO_NDC(uv);
	const F32 depth = textureLod(u_depthRt, u_linearAnyClampSampler, uv, 0.0).r;
	const Vec4 worldPos4 = u_clusteredShading.m_matrices.m_invertedViewProjectionJitter * Vec4(ndc, depth, 1.0);
	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;

	// Cluster
	const Vec2 fragCoord = uv * u_clusteredShading.m_renderingSize;
	Cluster cluster = getClusterFragCoord(Vec3(fragCoord, depth), TILE_SIZE, TILE_COUNTS, Z_SPLIT_COUNT,
										  u_clusteredShading.m_zSplitMagic.x, u_clusteredShading.m_zSplitMagic.y);

	// Layers
	U32 shadowCasterCountPerFragment = 0u;
	const U32 maxShadowCastersPerFragment = 4u;
	F32 shadowFactors[maxShadowCastersPerFragment] = F32[](0.0, 0.0, 0.0, 0.0);

	// Dir light
	const DirectionalLight dirLight = u_clusteredShading.m_directionalLight;
	if(dirLight.m_active != 0u && dirLight.m_cascadeCount > 0u)
	{
		const F32 positiveZViewSpace =
			testPlanePoint(u_clusteredShading.m_nearPlaneWSpace.xyz, u_clusteredShading.m_nearPlaneWSpace.w, worldPos)
			+ u_clusteredShading.m_near;

		F32 shadowFactor;
		if(positiveZViewSpace < dirLight.m_effectiveShadowDistance)
		{
			const U32 cascadeIdx =
				computeShadowCascadeIndex(positiveZViewSpace, dirLight.m_shadowCascadesDistancePower,
										  dirLight.m_effectiveShadowDistance, dirLight.m_cascadeCount);

			shadowFactor =
				computeShadowFactorDirLight(dirLight, cascadeIdx, worldPos, u_shadowAtlasTex, u_linearAnyClampSampler);

			F32 distanceFadeFactor = saturate(positiveZViewSpace / dirLight.m_effectiveShadowDistance);
			distanceFadeFactor = pow(distanceFadeFactor, 8.0);
			shadowFactor += distanceFadeFactor;
		}
		else
		{
			shadowFactor = 1.0;
		}

		shadowFactors[0] = shadowFactor;
		++shadowCasterCountPerFragment;
	}

	// Point lights
	ANKI_LOOP while(cluster.m_pointLightsMask != 0ul)
	{
		const I32 idx = findLSB64(cluster.m_pointLightsMask);
		cluster.m_pointLightsMask &= ~(1ul << U64(idx));
		const PointLight light = u_pointLights2[idx];

		ANKI_BRANCH if(light.m_shadowAtlasTileScale >= 0.0)
		{
			const Vec3 frag2Light = light.m_position - worldPos;

			const F32 shadowFactor =
				computeShadowFactorPointLight(light, frag2Light, u_shadowAtlasTex, u_linearAnyClampSampler);
			shadowFactors[min(maxShadowCastersPerFragment - 1u, shadowCasterCountPerFragment++)] = shadowFactor;
		}
	}

	// Spot lights
	ANKI_LOOP while(cluster.m_spotLightsMask != 0ul)
	{
		const I32 idx = findLSB64(cluster.m_spotLightsMask);
		cluster.m_spotLightsMask &= ~(1ul << U64(idx));
		const SpotLight light = u_spotLights2[idx];

		ANKI_BRANCH if(light.m_shadowLayer != MAX_U32)
		{
			const F32 shadowFactor =
				computeShadowFactorSpotLight(light, worldPos, u_shadowAtlasTex, u_linearAnyClampSampler);
			shadowFactors[min(maxShadowCastersPerFragment - 1u, shadowCasterCountPerFragment++)] = shadowFactor;
		}
	}

	// Store
	imageStore(u_outImg, IVec2(gl_GlobalInvocationID.xy),
			   Vec4(shadowFactors[0], shadowFactors[1], shadowFactors[2], shadowFactors[3]));
}

#pragma anki end