// 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 BLUR_ORIENTATION 0 1 // 0: in X asix, 1: in Y axis

#pragma anki start comp

#include <AnKi/Shaders/Include/IndirectDiffuseTypes.h>
#include <AnKi/Shaders/PackFunctions.glsl>
#include <AnKi/Shaders/Functions.glsl>
#include <AnKi/Shaders/BilateralFilter.glsl>

const UVec2 WORKGROUP_SIZE = UVec2(8u, 8u);
layout(local_size_x = WORKGROUP_SIZE.x, local_size_y = WORKGROUP_SIZE.y) in;

layout(set = 0, binding = 0) uniform sampler u_linearAnyClampSampler;
layout(set = 0, binding = 1) uniform texture2D u_toDenoiseTex;
layout(set = 0, binding = 2) uniform texture2D u_depthTex;
layout(set = 0, binding = 3) uniform texture2D u_gbuffer2Tex;
layout(set = 0, binding = 4) writeonly uniform image2D u_outImg;

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

Vec3 unproject(Vec2 ndc, F32 depth)
{
	const Vec4 worldPos4 = u_unis.m_invertedViewProjectionJitterMat * Vec4(ndc, depth, 1.0);
	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;
	return worldPos;
}

void main()
{
	if(skipOutOfBoundsInvocations(WORKGROUP_SIZE, u_unis.m_viewportSize))
	{
		return;
	}

	const Vec2 uv = (Vec2(gl_GlobalInvocationID.xy) + 0.5) / u_unis.m_viewportSizef;

	// Reference
	const F32 depthCenter = textureLod(u_depthTex, u_linearAnyClampSampler, uv, 0.0).r;
	if(depthCenter == 1.0)
	{
		imageStore(u_outImg, IVec2(gl_GlobalInvocationID.xy), Vec4(0.0));
		return;
	}

	const Vec3 positionCenter = unproject(UV_TO_NDC(uv), depthCenter);
	const Vec3 normalCenter = readNormalFromGBuffer(u_gbuffer2Tex, u_linearAnyClampSampler, uv);

	// Sample
	F32 weight = EPSILON;
	Vec3 color = Vec3(0.0);

	for(F32 i = -u_unis.m_sampleCountDiv2; i <= u_unis.m_sampleCountDiv2; i += 1.0)
	{
		const Vec2 texelSize = 1.0 / u_unis.m_viewportSizef;
#if BLUR_ORIENTATION == 0
		const Vec2 sampleUv = Vec2(uv.x + i * texelSize.x, uv.y);
#else
		const Vec2 sampleUv = Vec2(uv.x, uv.y + i * texelSize.y);
#endif

		const F32 depthTap = textureLod(u_depthTex, u_linearAnyClampSampler, sampleUv, 0.0).r;
		const Vec3 positionTap = unproject(UV_TO_NDC(sampleUv), depthTap);
		const Vec3 normalTap =
			unpackNormalFromGBuffer(textureLod(u_gbuffer2Tex, u_linearAnyClampSampler, sampleUv, 0.0));

		F32 w = calculateBilateralWeightPlane(positionCenter, normalCenter, positionTap, normalTap, 1.0);
		// w *= gaussianWeight(0.4, abs(F32(i)) / (sampleCount + 1.0));
		weight += w;

		color += textureLod(u_toDenoiseTex, u_linearAnyClampSampler, sampleUv, 0.0).xyz * w;
	}

	// Normalize and store
	color /= weight;
	imageStore(u_outImg, IVec2(gl_GlobalInvocationID.xy), Vec4(color, 0.0));
	// imageStore(u_outImg, IVec2(gl_GlobalInvocationID.xy), textureLod(u_toDenoiseTex, u_linearAnyClampSampler, uv,
	// 0.0));
}

#pragma anki end