// 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 RAYS_PER_PIXEL 1 2 4 8
#pragma anki library RtShadows
#pragma anki ray_type 0

#pragma anki start rgen

#include <AnKi/Shaders/ImportanceSampling.glsl>
#include <AnKi/Shaders/PackFunctions.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_CLUSTERS_BINDING 4
#include <AnKi/Shaders/ClusteredShadingCommon.glsl>

// Used by the hit shaders. When changing the binding you need to update other shaders
layout(set = 0, binding = 5) uniform sampler u_trilinearRepeatSampler;

layout(set = 0, binding = 6) uniform uimage2D u_shadowsImage;
layout(set = 0, binding = 7) uniform utexture2D u_historyShadowsTex;
layout(set = 0, binding = 8) uniform sampler u_linearAnyClampSampler;
layout(set = 0, binding = 9) uniform sampler u_nearestAnyClampSampler;
layout(set = 0, binding = 10) uniform texture2D u_depthRt;
layout(set = 0, binding = 11) uniform texture2D u_motionVectorsRt;
layout(set = 0, binding = 12) uniform texture2D u_motionVectorsRejectionRt;
layout(set = 0, binding = 13) uniform texture2D u_normalRt;
layout(set = 0, binding = 14) uniform accelerationStructureEXT u_tlas;
layout(set = 0, binding = 15) uniform texture2D u_prevHistoryLengthTex;
layout(set = 0, binding = 16) uniform image2D u_historyLengthImage;
layout(set = 0, binding = 17) uniform texture2D u_prevMomentsTex;
layout(set = 0, binding = 18) uniform image2D u_momentsImage;
layout(set = 0, binding = 19) uniform texture2D u_blueNoiseTex;

ANKI_BINDLESS_SET(1); // Used by the hit shaders

layout(push_constant, std430) uniform b_pc
{
	RtShadowsUniforms u_unis;
};

layout(location = 0) rayPayloadEXT F32 g_payload;

F32 trace(const Vec3 rayOrigin, const Vec3 rayDir, F32 tMax)
{
	const U32 flags = gl_RayFlagsOpaqueEXT | gl_RayFlagsTerminateOnFirstHitEXT | gl_RayFlagsSkipClosestHitShaderEXT;
	const U32 cullMask = 0xFFu;
	const U32 sbtRecordOffset = 0u;
	const U32 sbtRecordStride = 0u;
	const U32 missIndex = 0u;
	const F32 tMin = 0.1;
	const I32 payloadLocation = 0;
	g_payload = 0.0;
	traceRayEXT(u_tlas, flags, cullMask, sbtRecordOffset, sbtRecordStride, missIndex, rayOrigin, tMin, rayDir, tMax,
				payloadLocation);

	return g_payload;
}

void main()
{
	// World position
	const Vec2 uv = (Vec2(gl_LaunchIDEXT.xy) + 0.5) / Vec2(gl_LaunchSizeEXT.xy);
	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;

	if(depth == 1.0)
	{
		imageStore(u_shadowsImage, IVec2(gl_LaunchIDEXT.xy), UVec4(0));
		imageStore(u_momentsImage, IVec2(gl_LaunchIDEXT.xy), Vec4(0.0));

		// Set to max history length because this pixel won't need any processing from further compute stages
		imageStore(u_historyLengthImage, IVec2(gl_LaunchIDEXT.xy),
				   Vec4(RT_SHADOWS_MAX_HISTORY_LENGTH / RT_SHADOWS_MAX_HISTORY_LENGTH, 0.0, 0.0, 0.0));

		return;
	}

	// World normal
	const Vec3 normal = readNormalFromGBuffer(u_normalRt, u_linearAnyClampSampler, uv);

	// Cluster
	Cluster cluster = getClusterFragCoord(Vec3(uv * u_clusteredShading.m_renderingSize, depth));

	F32 shadowFactors[MAX_RT_SHADOW_LAYERS];
	zeroRtShadowLayers(shadowFactors);

	// Get a random factor
	Vec3 random[RAYS_PER_PIXEL];
	for(I32 i = 0; i < RAYS_PER_PIXEL; ++i)
	{
		const U32 frameIdx = u_clusteredShading.m_frame * U32(RAYS_PER_PIXEL + i);
#if 0
		const UVec3 irandom = rand3DPCG16(UVec3(gl_LaunchIDEXT.xy, frameIdx));
		random[i] = Vec3(irandom) / F32(0xFFFF) * 2.0 - 1.0; // In [-1.0, 1.0]
#else
		random[i] =
			textureLod(u_blueNoiseTex, u_trilinearRepeatSampler, Vec2(gl_LaunchSizeEXT.xy) / Vec2(64.0) * uv, 0.0).rgb;
		random[i] = animateBlueNoise(random[i], frameIdx);
		random[i] = random[i] * 2.0 - 1.0; // In [-1.0, 1.0]
#endif
	}

	// Dir light
	const DirectionalLight dirLight = u_clusteredShading.m_directionalLight;
	ANKI_BRANCH if(dirLight.m_active != 0u && dirLight.m_cascadeCount > 0u)
	{
		for(I32 i = 0; i < RAYS_PER_PIXEL; ++i)
		{
			const Vec3 dirLightPos = worldPos + -dirLight.m_direction * 10.0 + random[i];
			const Vec3 rayDir = normalize(dirLightPos - worldPos);

			const F32 lambertTerm = dot(rayDir, normal);
			ANKI_BRANCH if(lambertTerm > 0.0)
			{
				shadowFactors[dirLight.m_shadowLayer] += trace(worldPos, rayDir, 10000.0) / F32(RAYS_PER_PIXEL);
			}
		}
	}

	// 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];

		ANKI_BRANCH if(light.m_shadowAtlasTileScale >= 0.0)
		{
			for(I32 i = 0; i < RAYS_PER_PIXEL; ++i)
			{
				const Vec3 lightPos = light.m_position + 0.05 * light.m_radius * random[i];
				const Vec3 toLight = lightPos - worldPos;
				const F32 distanceToLight = length(toLight);
				const Vec3 rayDir = toLight / distanceToLight; // normalize

				const Bool inside = distanceToLight < light.m_radius;
				const F32 lambertTerm = dot(rayDir, normal);
				ANKI_BRANCH if(inside && lambertTerm > 0.0)
				{
					shadowFactors[light.m_shadowLayer] +=
						trace(worldPos, rayDir, distanceToLight) / F32(RAYS_PER_PIXEL);
				}
			}
		}
	}

	// 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];

		ANKI_BRANCH if(light.m_shadowLayer != MAX_U32)
		{
			for(I32 i = 0; i < RAYS_PER_PIXEL; ++i)
			{
				const Vec3 lightPos = light.m_position + 0.05 * light.m_radius * random[i];
				const Vec3 toLight = lightPos - worldPos;
				const F32 distanceToLight = length(toLight);
				const Vec3 rayDir = toLight / distanceToLight; // normalize

				const F32 lambertTerm = dot(rayDir, normal);
				ANKI_BRANCH if(lambertTerm > 0.0)
				{
					shadowFactors[light.m_shadowLayer] +=
						trace(worldPos, rayDir, distanceToLight) / F32(RAYS_PER_PIXEL);
				}
			}
		}
	}

	// Compute history length
	const Vec2 historyUv = uv + textureLod(u_motionVectorsRt, u_linearAnyClampSampler, uv, 0.0).xy;
	const F32 historyRejectionFactor = textureLod(u_motionVectorsRejectionRt, u_linearAnyClampSampler, uv, 0.0).x;
	F32 historyLength;
	if(historyRejectionFactor >= 0.5)
	{
		// Rejection factor too high, reset the temporal history for all layers
		historyLength = 1.0 / RT_SHADOWS_MAX_HISTORY_LENGTH;
	}
	else
	{
		// Sample seems stable, increment its temporal history

		historyLength = textureLod(u_prevHistoryLengthTex, u_linearAnyClampSampler, historyUv, 0.0).r;
		historyLength += 1.0 / RT_SHADOWS_MAX_HISTORY_LENGTH;
	}

	// Store history length
	imageStore(u_historyLengthImage, IVec2(gl_LaunchIDEXT.xy), Vec4(historyLength));

	// Compute blend fractor. Use nearest sampler because it's an integer texture
	const F32 lowestBlendFactor = 0.1;
	const F32 stableFrames = 4.0;
	const F32 lerp = min(1.0, (historyLength * RT_SHADOWS_MAX_HISTORY_LENGTH - 1.0) / stableFrames);
	const F32 blendFactor = mix(1.0, lowestBlendFactor, lerp);

	// Blend with history
	const UVec4 packedhistory = textureLod(u_historyShadowsTex, u_nearestAnyClampSampler, historyUv, 0.0);
	F32 history[MAX_RT_SHADOW_LAYERS];
	unpackRtShadows(packedhistory, history);
	for(U32 i = 0u; i < MAX_RT_SHADOW_LAYERS; ++i)
	{
		const F32 lerp = min(1.0, u_unis.historyRejectFactor[i] + blendFactor);
		shadowFactors[i] = mix(history[i], shadowFactors[i], lerp);
	}

	// Store the shadows image
	const UVec4 packed = packRtShadows(shadowFactors);
	imageStore(u_shadowsImage, IVec2(gl_LaunchIDEXT.xy), packed);

	// Compute the moments that will give temporal variance
	Vec2 moments = Vec2(0.0);
	ANKI_UNROLL for(U32 i = 0u; i < MAX_RT_SHADOW_LAYERS; ++i)
	{
		moments.x += shadowFactors[i];
	}
	moments.y = moments.x * moments.x;

	// Blend the moments
	const Vec2 prevMoments = textureLod(u_prevMomentsTex, u_linearAnyClampSampler, historyUv, 0.0).xy;
	F32 momentsBlendFactor = 0.2;
	momentsBlendFactor = mix(momentsBlendFactor, 1.0, historyRejectionFactor);
	moments = mix(prevMoments, moments, momentsBlendFactor);

	// Store the moments
	imageStore(u_momentsImage, IVec2(gl_LaunchIDEXT.xy), Vec4(moments, 0.0, 0.0));
}
#pragma anki end