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

#pragma anki mutator QUALITY 0 1 2 // No filtering, PCF, PCSS
#pragma anki mutator DIRECTIONAL_LIGHT_SHADOW_RESOLVED 0 1

#pragma anki technique vert pixel comp

#include <AnKi/Shaders/QuadVert.hlsl>
#if ANKI_COMPUTE_SHADER || ANKI_PIXEL_SHADER
#	include <AnKi/Shaders/ClusteredShadingFunctions.hlsl>
#	include <AnKi/Shaders/ImportanceSampling.hlsl>

#	define DEBUG_CASCADES 0

ConstantBuffer<GlobalRendererConstants> g_globalConstants : register(b0);
StructuredBuffer<PointLight> g_pointLights : register(t0);
StructuredBuffer<SpotLight> g_spotLights : register(t1);
Texture2D<Vec4> g_shadowAtlasTex : register(t2);
StructuredBuffer<Cluster> g_clusters : register(t3);

SamplerState g_linearAnyClampSampler : register(s0);
SamplerComparisonState g_linearAnyClampShadowSampler : register(s1);
SamplerState g_trilinearRepeatSampler : register(s2);
Texture2D<Vec4> g_depthRt : register(t4);
Texture2D<Vec4> g_noiseTex : register(t5);

#	if DIRECTIONAL_LIGHT_SHADOW_RESOLVED
Texture2D<Vec4> g_dirLightResolvedShadowsTex : register(t6);
#	endif

#	if ANKI_COMPUTE_SHADER
RWTexture2D<Vec4> g_storageTex : register(u0);
#	endif

struct Constants
{
	Vec2 m_framebufferSize;
	F32 m_padding0;
	F32 m_padding1;
};

ANKI_FAST_CONSTANTS(Constants, g_consts)

Vec3 computeDebugShadowCascadeColor(U32 cascade)
{
	if(cascade == 0u)
	{
		return Vec3(0.0f, 1.0f, 0.0f);
	}
	else if(cascade == 1u)
	{
		return Vec3(0.0f, 0.0f, 1.0f);
	}
	else if(cascade == 2u)
	{
		return Vec3(0.0f, 1.0f, 1.0f);
	}
	else
	{
		return Vec3(1.0f, 0.0f, 0.0f);
	}
}

#	if ANKI_COMPUTE_SHADER
[numthreads(8, 8, 1)] void main(UVec2 svDispatchThreadId : SV_DISPATCHTHREADID)
#	else
Vec4 main(VertOut input) : SV_TARGET0
#	endif
{
#	if ANKI_COMPUTE_SHADER
	svDispatchThreadId = min(svDispatchThreadId, UVec2(g_consts.m_framebufferSize - 1.0f)); // Just to be sure
	const Vec2 uv = (Vec2(svDispatchThreadId) + 0.5) / g_consts.m_framebufferSize;
#	else
	const Vec2 uv = input.m_uv;
	const UVec2 svDispatchThreadId = input.m_svPosition;
	ANKI_MAYBE_UNUSED(svDispatchThreadId);
#	endif

#	if QUALITY > 0
#		if 1
	// Noise
	Vec2 noiseTexSize;
	g_noiseTex.GetDimensions(noiseTexSize.x, noiseTexSize.y);

	const Vec2 noiseUv = g_consts.m_framebufferSize / noiseTexSize * uv;
	Vec3 noise = g_noiseTex.SampleLevel(g_trilinearRepeatSampler, noiseUv, 0.0).rgb;
	noise = animateBlueNoise(noise, g_globalConstants.m_frame % 16u);
	const F32 randFactor = noise.x;
#		else
	const Vec2 noise2 = spatioTemporalNoise(svDispatchThreadId, g_globalConstants.m_frame);
	const F32 randFactor = noise2.x;
#		endif
#	endif

	// World position
	const Vec2 ndc = uvToNdc(uv);
	const F32 depth = g_depthRt.SampleLevel(g_linearAnyClampSampler, uv, 0.0).r;
	const Vec4 worldPos4 = mul(g_globalConstants.m_matrices.m_invertedViewProjectionJitter, Vec4(ndc, depth, 1.0));
	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;

	// Cluster
	const Vec2 fragCoord = uv * g_globalConstants.m_renderingSize;
	Cluster cluster = getClusterFragCoord(g_clusters, g_globalConstants, Vec3(fragCoord, depth));

	// Layers
	U32 shadowCasterCountPerFragment = 0u;
	const U32 kMaxShadowCastersPerFragment = 4u;
	Vec4 shadowFactors = 0.0f;

	// Dir light
#	if DIRECTIONAL_LIGHT_SHADOW_RESOLVED
	shadowFactors[0] = g_dirLightResolvedShadowsTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0f).x;
	++shadowCasterCountPerFragment;
#	else
	const DirectionalLight dirLight = g_globalConstants.m_directionalLight;
	if(dirLight.m_active && dirLight.m_shadowCascadeCount)
	{
		const U32 shadowCascadeCount = dirLight.m_shadowCascadeCount;

		const F32 positiveZViewSpace = testPlanePoint(g_globalConstants.m_nearPlaneWSpace.xyz, g_globalConstants.m_nearPlaneWSpace.w, worldPos)
									   + g_globalConstants.m_matrices.m_near;

		const F32 lastCascadeDistance = dirLight.m_shadowCascadeDistances[shadowCascadeCount - 1u];
		F32 shadowFactor;
		if(positiveZViewSpace < lastCascadeDistance)
		{
			F32 cascadeBlendFactor;
			const UVec2 cascadeIndices =
				computeShadowCascadeIndex2(positiveZViewSpace, dirLight.m_shadowCascadeDistances, shadowCascadeCount, cascadeBlendFactor);

#		if DEBUG_CASCADES
			const Vec3 debugColorA = computeDebugShadowCascadeColor(cascadeIndices[0]);
			const Vec3 debugColorB = computeDebugShadowCascadeColor(cascadeIndices[1]);
			const Vec3 debugColor = lerp(debugColorA, debugColorB, cascadeBlendFactor);
#			if ANKI_COMPUTE_SHADER
			g_storageTex[svDispatchThreadId.xy] = shadowFactors;
			return;
#			else
			return shadowFactors;
#			endif
#		endif

#		if QUALITY == 2
			const F32 shadowFactorCascadeA = computeShadowFactorDirLightPcss<F32>(dirLight, cascadeIndices.x, worldPos, g_shadowAtlasTex,
																				  g_linearAnyClampShadowSampler, randFactor, g_linearAnyClampSampler);
#		elif QUALITY == 1
			const F32 shadowFactorCascadeA = computeShadowFactorDirLightPcf<F32>(dirLight, cascadeIndices.x, worldPos, g_shadowAtlasTex,
																				 g_linearAnyClampShadowSampler, randFactor);
#		else
			const F32 shadowFactorCascadeA =
				computeShadowFactorDirLight<F32>(dirLight, cascadeIndices.x, worldPos, g_shadowAtlasTex, g_linearAnyClampShadowSampler);
#		endif

			if(cascadeBlendFactor < 0.01 || cascadeIndices.x == cascadeIndices.y)
			{
				// Don't blend cascades
				shadowFactor = shadowFactorCascadeA;
			}
			else
			{
#		if QUALITY == 2
				// Blend cascades
				const F32 shadowFactorCascadeB = computeShadowFactorDirLightPcss<F32>(
					dirLight, cascadeIndices.y, worldPos, g_shadowAtlasTex, g_linearAnyClampShadowSampler, randFactor, g_linearAnyClampSampler);
#		elif QUALITY == 1
				// Blend cascades
				const F32 shadowFactorCascadeB = computeShadowFactorDirLightPcf<F32>(dirLight, cascadeIndices.y, worldPos, g_shadowAtlasTex,
																					 g_linearAnyClampShadowSampler, randFactor);
#		else
				// Blend cascades
				const F32 shadowFactorCascadeB =
					computeShadowFactorDirLight<F32>(dirLight, cascadeIndices.y, worldPos, g_shadowAtlasTex, g_linearAnyClampShadowSampler);
#		endif
				shadowFactor = lerp(shadowFactorCascadeA, shadowFactorCascadeB, cascadeBlendFactor);
			}

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

		shadowFactors[0] = shadowFactor;
		++shadowCasterCountPerFragment;
	}
#	endif // DIRECTIONAL_LIGHT_SHADOW_RESOLVED

	// Point lights
	U32 idx = 0;
	[loop] while((idx = iteratePointLights(cluster)) != kMaxU32)
	{
		const PointLight light = g_pointLights[idx];

		[branch] if(light.m_shadow)
		{
			const Vec3 frag2Light = light.m_position - worldPos;

#	if QUALITY > 0
			const F32 shadowFactor = computeShadowFactorPointLightPcf(light, frag2Light, g_shadowAtlasTex, g_linearAnyClampShadowSampler, randFactor);
#	else
			const F32 shadowFactor = computeShadowFactorPointLight<F32>(light, frag2Light, g_shadowAtlasTex, g_linearAnyClampShadowSampler);
#	endif
			shadowFactors[min(kMaxShadowCastersPerFragment - 1u, shadowCasterCountPerFragment++)] = shadowFactor;
		}
	}

	// Spot lights
	[loop] while((idx = iterateSpotLights(cluster)) != kMaxU32)
	{
		const SpotLight light = g_spotLights[idx];

		[branch] if(light.m_shadow)
		{
#	if QUALITY == 2
			const F32 shadowFactor = computeShadowFactorSpotLightPcss<F32>(light, worldPos, g_shadowAtlasTex, g_linearAnyClampShadowSampler,
																		   randFactor, g_linearAnyClampSampler);
#	elif QUALITY == 1
			const F32 shadowFactor =
				computeShadowFactorSpotLightPcf<F32>(light, worldPos, g_shadowAtlasTex, g_linearAnyClampShadowSampler, randFactor);
#	else
			const F32 shadowFactor = computeShadowFactorSpotLight<F32>(light, worldPos, g_shadowAtlasTex, g_linearAnyClampShadowSampler);
#	endif
			shadowFactors[min(kMaxShadowCastersPerFragment - 1u, shadowCasterCountPerFragment++)] = shadowFactor;
		}
	}

	// Store
#	if ANKI_COMPUTE_SHADER
	g_storageTex[svDispatchThreadId] = shadowFactors;
#	else
	return shadowFactors;
#	endif
}
#endif // ANKI_COMPUTE_SHADER || ANKI_PIXEL_SHADER