// 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 BLUR_ORIENTATION 0 1

#pragma anki technique comp

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

SamplerState g_linearAnyClampSampler : register(s0);
Texture2D<Vec4> g_inTex : register(t0);
Texture2D<Vec4> g_depthTex : register(t1);
Texture2D<Vec4> g_gbuffer2Tex : register(t2);
Texture2D<Vec4> g_momentsTex : register(t3);
Texture2D<Vec4> g_historyLengthTex : register(t4);
RWTexture2D<Vec4> g_storageTex : register(u0);

ANKI_FAST_CONSTANTS(RtShadowsDenoiseConstants, g_consts)

Vec3 unproject(Vec2 ndc, F32 depth)
{
	const Vec4 worldPos4 = mul(g_consts.m_invViewProjMat, Vec4(ndc, depth, 1.0));
	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;
	return worldPos;
}

F32 computeVarianceCenter(Vec2 uv)
{
	const F32 kernel[2][2] = {{1.0 / 4.0, 1.0 / 8.0}, {1.0 / 8.0, 1.0 / 16.0}};
	const I32 radius = 1;
	Vec2 texSize;
	F32 mipCount;
	g_momentsTex.GetDimensions(0, texSize.x, texSize.y, mipCount);
	const Vec2 texelSize = 1.0 / texSize;
	Vec2 sumMoments = 0.0f;

	for(I32 yy = -radius; yy <= radius; yy++)
	{
		for(I32 xx = -radius; xx <= radius; xx++)
		{
			const Vec2 newUv = uv + Vec2(xx, yy) * texelSize;
			const F32 k = kernel[abs(xx)][abs(yy)];
			sumMoments += g_momentsTex.SampleLevel(g_linearAnyClampSampler, newUv, 0.0).xy * k;
		}
	}

	return abs(sumMoments.y - sumMoments.x * sumMoments.x);
}

[numthreads(8, 8, 1)] void main(UVec3 svDispatchThreadId : SV_DISPATCHTHREADID)
{
	Vec2 outImageSize;
	g_storageTex.GetDimensions(outImageSize.x, outImageSize.y);
	if(any(svDispatchThreadId.xy >= UVec2(outImageSize)))
	{
		return;
	}

	const Vec2 uv = (Vec2(svDispatchThreadId.xy) + 0.5) / outImageSize;

	// Reference
	const F32 depthCenter = g_depthTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).r;
	if(depthCenter == 1.0)
	{
		g_storageTex[svDispatchThreadId.xy] = 0.0;
		return;
	}

	const Vec3 positionCenter = unproject(uvToNdc(uv), depthCenter);
	const Vec3 normalCenter = unpackNormalFromGBuffer(g_gbuffer2Tex.SampleLevel(g_linearAnyClampSampler, uv, 0.0));

	// Decide the amount of blurring
	const F32 varianceCenter = computeVarianceCenter(uv);
	const F32 historyLength = g_historyLengthTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).x * kRtShadowsMaxHistoryLength;

	U32 sampleCount;
	if(historyLength < 2.0)
	{
		// Worst case
		sampleCount = g_consts.m_maxSampleCount;
	}
	else if(historyLength > 4.0 && varianceCenter < 0.0001)
	{
		// Best case
		sampleCount = g_consts.m_minSampleCount;
	}
	else
	{
		// Every other case

		F32 blur = varianceCenter * 100.0;
		blur = min(1.0, blur);

		const F32 sampleCountf = lerp(F32(g_consts.m_minSampleCount), F32(g_consts.m_maxSampleCount), blur);

		sampleCount = U32(sampleCountf);
	}

	sampleCount = sampleCount / 2u;

	// Sample
	F32 weight = kEpsilonF32;
	F32 shadowFactor = 0.0;
	for(I32 i = -I32(sampleCount); i < I32(sampleCount); ++i)
	{
		const Vec2 texelSize = 1.0 / outImageSize;
#if BLUR_ORIENTATION == 0
		const Vec2 sampleUv = Vec2(uv.x + F32(i) * texelSize.x, uv.y);
#else
		const Vec2 sampleUv = Vec2(uv.x, uv.y + F32(i) * texelSize.y);
#endif

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

		const F32 depthTap = g_depthTex.SampleLevel(g_linearAnyClampSampler, sampleUv, 0.0).r;
		const Vec3 positionTap = unproject(uvToNdc(sampleUv), depthTap);
		const Vec3 normalTap = unpackNormalFromGBuffer(g_gbuffer2Tex.SampleLevel(g_linearAnyClampSampler, sampleUv, 0.0));

		// F32 localWeight = calculateBilateralWeighPlane(depthCenter, depthTap, 1.0);
		F32 localWeight = calculateBilateralWeightPlane(positionCenter, normalCenter, positionTap, normalTap, 1.0);

		localWeight *= gaussianWeight(0.4f, abs(F32(i)) / F32(sampleCount + 1u));

		shadowFactor += localShadowFactor * localWeight;

		weight += localWeight;
	}

	// Normalize and write
	shadowFactor /= weight;
	g_storageTex[svDispatchThreadId.xy] = shadowFactor;
}