// 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 RAYS_PER_PIXEL 1 2 4 8

#pragma anki technique RtShadows rgen miss

#include <AnKi/Shaders/ImportanceSampling.hlsl>
#include <AnKi/Shaders/PackFunctions.hlsl>
#include <AnKi/Shaders/RtShadows.hlsl>
#include <AnKi/Shaders/MaterialShadersCommon.hlsl>
#include <AnKi/Shaders/ClusteredShadingFunctions.hlsl>

#if ANKI_RAY_GEN_SHADER
#	define SPACE space2

ConstantBuffer<GlobalRendererConstants> g_globalRendererConstants : register(b0, SPACE);

SamplerState g_trilinearRepeatSampler : register(s0, SPACE);

RWTexture2D<Vec4> g_shadowsImage : register(u0, SPACE);
Texture2D<Vec4> g_historyShadowsTex : register(t0, SPACE);
SamplerState g_linearAnyClampSampler : register(s1, SPACE);
Texture2D<Vec4> g_depthRt : register(t1, SPACE);
Texture2D<Vec4> g_motionVectorsRt : register(t2, SPACE);
Texture2D<Vec4> g_historyLengthTex : register(t3, SPACE);
Texture2D<Vec4> g_normalRt : register(t4, SPACE);
RaytracingAccelerationStructure g_tlas : register(t5, SPACE);
Texture2D<Vec4> g_prevMomentsTex : register(t6, SPACE);
RWTexture2D<Vec4> g_momentsImage : register(u1, SPACE);
Texture2D<Vec4> g_blueNoiseTex : register(t7, SPACE);

F32 trace(const Vec3 rayOrigin, const Vec3 rayDir, F32 tMax)
{
	const U32 flags = RAY_FLAG_ACCEPT_FIRST_HIT_AND_END_SEARCH | RAY_FLAG_SKIP_CLOSEST_HIT_SHADER;
	const U32 cullMask = 0xFFu;
	const U32 sbtRecordOffset = 0u;
	const U32 sbtRecordStride = 0u;
	const U32 missIndex = 0u;
	RayDesc ray;
	ray.Origin = rayOrigin;
	ray.TMin = 0.1;
	ray.Direction = rayDir;
	ray.TMax = tMax;

	RtShadowsRayPayload payload;
	payload.m_shadowFactor = 0.0;
	TraceRay(g_tlas, flags, cullMask, sbtRecordOffset, sbtRecordStride, missIndex, ray, payload);

	return payload.m_shadowFactor;
}

Vec3 genRandomDirection(U32 rayIdx, Vec2 uv)
{
	const U32 frameIdx = g_globalRendererConstants.m_frame * RAYS_PER_PIXEL + rayIdx;

	Vec2 noiseTexSize;
	g_blueNoiseTex.GetDimensions(noiseTexSize.x, noiseTexSize.y);

	Vec3 random = g_blueNoiseTex.SampleLevel(g_trilinearRepeatSampler, Vec2(DispatchRaysDimensions().xy) / noiseTexSize * uv, 0.0).rgb;
	random = animateBlueNoise(random, frameIdx);
	random = random * 2.0 - 1.0; // In [-1.0, 1.0]

	return random;
}

[shader("raygeneration")] void main()
{
	// World position
	const Vec2 uv = (Vec2(DispatchRaysIndex().xy) + 0.5) / Vec2(DispatchRaysDimensions().xy);
	const Vec2 ndc = uvToNdc(uv);
	const F32 depth = g_depthRt.SampleLevel(g_linearAnyClampSampler, uv, 0.0).r;
	const Vec4 worldPos4 = mul(g_globalRendererConstants.m_matrices.m_invertedViewProjectionJitter, Vec4(ndc, depth, 1.0));
	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;

	if(depth == 1.0)
	{
		g_shadowsImage[DispatchRaysIndex().xy] = 0.0f;
		g_momentsImage[DispatchRaysIndex().xy] = 0.0f;
		return;
	}

	// World normal
	const Vec3 normal = unpackNormalFromGBuffer(g_normalRt.SampleLevel(g_linearAnyClampSampler, uv, 0.0));

	// Dir light
	F32 shadowFactor = 0.0f;
	const DirectionalLight dirLight = g_globalRendererConstants.m_directionalLight;
	for(U32 i = 0; i < RAYS_PER_PIXEL; ++i)
	{
		const Vec3 dirLightPos = worldPos + -dirLight.m_direction * 10.0 + genRandomDirection(i, uv);
		const Vec3 rayDir = normalize(dirLightPos - worldPos);

		const F32 lambertTerm = dot(rayDir, normal);
		[branch] if(lambertTerm > 0.0)
		{
			shadowFactor += trace(worldPos, rayDir, 10000.0) / F32(RAYS_PER_PIXEL);
		}
	}

	// Get history length
	const Vec2 historyUv = uv + g_motionVectorsRt.SampleLevel(g_linearAnyClampSampler, uv, 0.0).xy;
	const F32 historyLength = g_historyLengthTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).x * kRtShadowsMaxHistoryLength;

	// Compute blend fractor
	const F32 lowestBlendFactor = 0.1;
	const F32 stableFrames = 4.0;
	const F32 historyGoodnessFactor = min(1.0, (historyLength - 1.0) / stableFrames);
	const F32 historyBlendFactor = lerp(1.0, lowestBlendFactor, historyGoodnessFactor);

	// Blend with history
	const F32 history = g_historyShadowsTex.SampleLevel(g_linearAnyClampSampler, historyUv, 0.0).r;
	shadowFactor = lerp(history, shadowFactor, historyBlendFactor);

	// Store the shadows image
	g_shadowsImage[DispatchRaysIndex().xy] = shadowFactor;

	// Compute the moments that will give temporal variance
	Vec2 moments;
	moments.x = shadowFactor;
	moments.y = moments.x * moments.x;

	// Blend the moments
	const Vec2 prevMoments = g_prevMomentsTex.SampleLevel(g_linearAnyClampSampler, historyUv, 0.0).xy;
	const F32 lowestMomentsBlendFactor = 0.2;
	const F32 momentsBlendFactor = lerp(1.0, lowestMomentsBlendFactor, historyGoodnessFactor);
	moments = lerp(prevMoments, moments, momentsBlendFactor);

	// Store the moments
	g_momentsImage[DispatchRaysIndex().xy] = Vec4(moments, 0.0, 0.0);
}
#endif // ANKI_RAY_GEN_SHADER

#if ANKI_MISS_SHADER
[shader("miss")] void main(inout RtShadowsRayPayload payload)
{
	payload.m_shadowFactor = 1.0;
}
#endif // ANKI_MISS_SHADER