// 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 technique comp

#include <AnKi/Shaders/RtShadows.hlsl>
#include <AnKi/Shaders/BilateralFilter.hlsl>
#include <AnKi/Shaders/Functions.hlsl>

SamplerState g_linearAnyClampSampler : register(s0);
Texture2D<Vec4> g_shadowsTex : register(t0);
Texture2D<Vec4> g_momentsTex : register(t1);
Texture2D<Vec4> g_historyLengthTex : register(t2);
Texture2D<Vec4> g_depthTex : register(t3);

RWTexture2D<Vec4> g_shadowUav : register(u0);
RWTexture2D<Vec4> g_varianceUav : register(u1);

ANKI_FAST_CONSTANTS(Mat4, g_invProjMat)

constexpr I32 kConvolutionRadius = 1;

Vec3 toViewspace(Vec2 uv, F32 depth)
{
	const Vec4 pos4 = mul(g_invProjMat, Vec4(uvToNdc(uv), depth, 1.0));
	const Vec3 pos = pos4.xyz / pos4.w;
	return pos;
}

[numthreads(8, 8, 1)] void main(UVec2 svDispatchThreadId : SV_DISPATCHTHREADID)
{
	Vec2 shadowUavSize;
	g_shadowUav.GetDimensions(shadowUavSize.x, shadowUavSize.y);

	if(any(Vec2(svDispatchThreadId) >= shadowUavSize))
	{
		return;
	}

	const Vec2 uv = (Vec2(svDispatchThreadId) + 0.5) / shadowUavSize;

	const F32 depth = g_depthTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).r;
	if(depth == 1.0)
	{
		// Sky
		g_shadowUav[svDispatchThreadId] = 0.0f;
		g_varianceUav[svDispatchThreadId] = 0.0f;
		return;
	}

	const F32 historyLength = g_historyLengthTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).r * kRtShadowsMaxHistoryLength;

	F32 outShadowFactor;
	F32 outVariance;

	if(historyLength < 4.0)
	{
		// It's been stable less than 4 frames, need to do some work

		const Vec2 texelSize = 1.0 / shadowUavSize;

		// Set the reference sample
		const F32 depthCenter = depth;
		const Vec3 positionCenter = toViewspace(uv, depthCenter);

		// Init the sums
		Vec2 sumMoments = 0.0f;
		F32 sumWeight = kEpsilonF32;
		F32 sumShadowFactor = 0.0f;

		// Convolve
		for(I32 offsetx = -kConvolutionRadius; offsetx <= kConvolutionRadius; offsetx++)
		{
			for(I32 offsety = -kConvolutionRadius; offsety <= kConvolutionRadius; offsety++)
			{
				const Vec2 sampleUv = uv + Vec2(offsetx, offsety) * texelSize;

				// Set the current sample
				const F32 depthTap = g_depthTex.SampleLevel(g_linearAnyClampSampler, sampleUv, 0.0).r;
				const Vec3 positionTap = toViewspace(sampleUv, depthTap);

				// Do bilateral
				const F32 w = calculateBilateralWeightViewspacePosition(positionCenter, positionTap, 0.5);

				// Sum
				const Vec2 moments = g_momentsTex.SampleLevel(g_linearAnyClampSampler, sampleUv, 0.0).xy;
				sumMoments += moments * w;

				const F32 localShadowFactor = g_shadowsTex.SampleLevel(g_linearAnyClampSampler, sampleUv, 0.0).r;

				sumShadowFactor += localShadowFactor * w;
				sumWeight += w;
			}
		}

		sumShadowFactor /= sumWeight;
		sumMoments /= sumWeight;

		outVariance = max(0.0, sumMoments.y - sumMoments.x * sumMoments.x);
		outShadowFactor = sumShadowFactor;

		// Give the variance a boost for the first frames
		outVariance *= 4.0 / (historyLength * kRtShadowsMaxHistoryLength);
	}
	else
	{
		// Stable for more that 4 frames, passthrough

		outShadowFactor = g_shadowsTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).r;

		const Vec2 moments = g_momentsTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).xy;
		outVariance = max(0.0, moments.y - moments.x * moments.x);
	}

	// Store
	g_shadowUav[svDispatchThreadId] = outShadowFactor;
	g_varianceUav[svDispatchThreadId] = Vec4(outVariance, 0.0, 0.0, 0.0);
}