// 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 mutator USE_SHADOW_LAYERS 0 1

ANKI_SPECIALIZATION_CONSTANT_UVEC2(TILE_COUNTS, 0u);
ANKI_SPECIALIZATION_CONSTANT_U32(Z_SPLIT_COUNT, 2u);
ANKI_SPECIALIZATION_CONSTANT_U32(TILE_SIZE, 3u);
ANKI_SPECIALIZATION_CONSTANT_U32(IR_MIPMAP_COUNT, 4u);

#pragma anki start vert
#include <AnKi/Shaders/QuadVert.glsl>
#pragma anki end

#pragma anki start frag
#include <AnKi/Shaders/PackFunctions.glsl>
#include <AnKi/Shaders/Functions.glsl>
#include <AnKi/Shaders/RtShadows.glsl>

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

layout(set = 0, binding = 8) uniform sampler u_nearestAnyClampSampler;
layout(set = 0, binding = 9) uniform sampler u_trilinearClampSampler;

layout(set = 0, binding = 10) uniform texture2D u_msRt0;
layout(set = 0, binding = 11) uniform texture2D u_msRt1;
layout(set = 0, binding = 12) uniform texture2D u_msRt2;
layout(set = 0, binding = 13) uniform texture2D u_msDepthRt;
layout(set = 0, binding = 14) uniform texture2D u_ssrRt;
#if USE_SHADOW_LAYERS
layout(set = 0, binding = 15) uniform utexture2D u_shadowLayersTex;
#else
layout(set = 0, binding = 16) uniform texture2D u_resolvedSm;
#endif

layout(location = 0) in Vec2 in_uv;

layout(location = 0) out Vec3 out_color;

// Common code for lighting
#define LIGHTING_COMMON_BRDF() \
	const Vec3 frag2Light = light.m_position - worldPos; \
	const Vec3 l = normalize(frag2Light); \
	const Vec3 specC = computeSpecularColorBrdf(gbuffer, viewDir, l); \
	const Vec3 diffC = diffuseLambert(gbuffer.m_diffuse); \
	const F32 att = computeAttenuationFactor(light.m_squareRadiusOverOne, frag2Light); \
	F32 lambert = max(0.0, dot(gbuffer.m_normal, l));

void main()
{
	const F32 depth = textureLod(u_msDepthRt, u_nearestAnyClampSampler, in_uv, 0.0).r;
	const Vec2 ndc = UV_TO_NDC(in_uv);

	if(depth == 1.0)
	{
		out_color = Vec3(0.0);
		return;
	}

	// Get world position
	const Vec4 worldPos4 = u_clusteredShading.m_matrices.m_invertedViewProjectionJitter * Vec4(ndc, depth, 1.0);
	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;

	// Get the cluster
	Cluster cluster = getClusterFragCoord(Vec3(gl_FragCoord.xy, depth), TILE_SIZE, TILE_COUNTS, Z_SPLIT_COUNT,
										  u_clusteredShading.m_zSplitMagic.x, u_clusteredShading.m_zSplitMagic.y);

	// out_color = clusterHeatmap(cluster, 1u << CLUSTER_OBJECT_TYPE_POINT_LIGHT); return;

	// Decode GBuffer
	GbufferInfo gbuffer;
	readGBuffer(u_msRt0, u_msRt1, u_msRt2, u_nearestAnyClampSampler, in_uv, 0.0, gbuffer);
	gbuffer.m_subsurface = max(gbuffer.m_subsurface, SUBSURFACE_MIN);

	// SM
#if USE_SHADOW_LAYERS
	F32 resolvedSm[MAX_RT_SHADOW_LAYERS];
	unpackRtShadows(textureLod(u_shadowLayersTex, u_nearestAnyClampSampler, in_uv, 0.0), resolvedSm);
#else
	Vec4 resolvedSm = textureLod(u_resolvedSm, u_trilinearClampSampler, in_uv, 0.0);
	U32 resolvedSmIdx = 0u;
#endif

	// Ambient and emissive color
	out_color = gbuffer.m_diffuse * gbuffer.m_emission;

	// Dir light
	const Vec3 viewDir = normalize(u_clusteredShading.m_cameraPosition - worldPos);
	const DirectionalLight dirLight = u_clusteredShading.m_directionalLight;
	if(dirLight.m_active != 0u)
	{
		F32 shadowFactor;
		if(dirLight.m_cascadeCount > 0u)
		{
#if USE_SHADOW_LAYERS
			shadowFactor = resolvedSm[dirLight.m_shadowLayer];
#else
			shadowFactor = resolvedSm[0];
			++resolvedSmIdx;
#endif
		}
		else
		{
			shadowFactor = 1.0;
		}

		const Vec3 l = -dirLight.m_direction;

		const F32 lambert = max(gbuffer.m_subsurface, dot(l, gbuffer.m_normal));

		const Vec3 diffC = diffuseLambert(gbuffer.m_diffuse);
		const Vec3 specC = computeSpecularColorBrdf(gbuffer, viewDir, l);

		out_color += (diffC + specC) * dirLight.m_diffuseColor * (shadowFactor * lambert);
	}

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

		LIGHTING_COMMON_BRDF();

		ANKI_BRANCH if(light.m_shadowAtlasTileScale >= 0.0)
		{
#if USE_SHADOW_LAYERS
			const F32 shadow = resolvedSm[light.m_shadowLayer];
#else
			const F32 shadow = resolvedSm[resolvedSmIdx++];
#endif
			lambert *= shadow;
		}

		out_color += (diffC + specC) * light.m_diffuseColor * (att * max(gbuffer.m_subsurface, lambert));
	}

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

		LIGHTING_COMMON_BRDF();

		const F32 spot = computeSpotFactor(l, light.m_outerCos, light.m_innerCos, light.m_direction);

		ANKI_BRANCH if(light.m_shadowLayer != MAX_U32)
		{
#if USE_SHADOW_LAYERS
			const F32 shadow = resolvedSm[light.m_shadowLayer];
#else
			const F32 shadow = resolvedSm[resolvedSmIdx++];
#endif
			lambert *= shadow;
		}

		out_color += (diffC + specC) * light.m_diffuseColor * (att * spot * max(gbuffer.m_subsurface, lambert));
	}

	// Indirect specular
	{
		// Do the probe read
		Vec3 specIndirect = Vec3(0.0);

		const Vec3 reflDir = reflect(-viewDir, gbuffer.m_normal);
		const F32 reflLod = F32(IR_MIPMAP_COUNT - 1u) * gbuffer.m_roughness;

		if(bitCount(cluster.m_reflectionProbesMask) == 1)
		{
			// Only one probe, do a fast path without blend weight

			const ReflectionProbe probe = u_reflectionProbes[findLSB2(cluster.m_reflectionProbesMask)];

			// Sample
			const Vec3 cubeUv = intersectProbe(worldPos, reflDir, probe.m_aabbMin, probe.m_aabbMax, probe.m_position);
			const Vec4 cubeArrUv = Vec4(cubeUv, probe.m_cubemapIndex);
			specIndirect = textureLod(u_reflectionsTex, u_trilinearClampSampler, cubeArrUv, reflLod).rgb;
		}
		else
		{
			// Zero or more than one probes, do a slow path that blends them together

			F32 totalBlendWeight = EPSILON;

			// Loop probes
			ANKI_LOOP while(cluster.m_reflectionProbesMask != 0u)
			{
				const U32 idx = U32(findLSB2(cluster.m_reflectionProbesMask));
				cluster.m_reflectionProbesMask &= ~(1u << idx);
				const ReflectionProbe probe = u_reflectionProbes[idx];

				// Compute blend weight
				const F32 blendWeight = computeProbeBlendWeight(worldPos, probe.m_aabbMin, probe.m_aabbMax, 0.2);
				totalBlendWeight += blendWeight;

				// Sample reflections
				const Vec3 cubeUv =
					intersectProbe(worldPos, reflDir, probe.m_aabbMin, probe.m_aabbMax, probe.m_position);
				const Vec4 cubeArrUv = Vec4(cubeUv, probe.m_cubemapIndex);
				const Vec3 c = textureLod(u_reflectionsTex, u_trilinearClampSampler, cubeArrUv, reflLod).rgb;
				specIndirect += c * blendWeight;
			}

			// Normalize the colors
			specIndirect /= totalBlendWeight;
		}

		// Read the SSL result
		const Vec4 ssr = textureLod(u_ssrRt, u_trilinearClampSampler, in_uv, 0.0);

		// Combine the SSR and probe reflections and write the result
		const Vec3 finalSpecIndirect = specIndirect * ssr.a + ssr.rgb;

		// Compute env BRDF
		const F32 NoV = max(EPSILON, dot(gbuffer.m_normal, viewDir));
		const Vec3 env =
			envBRDF(gbuffer.m_specular, gbuffer.m_roughness, u_integrationLut, u_trilinearClampSampler, NoV);

		out_color += finalSpecIndirect * env;
	}
}
#pragma anki end