// Copyright (C) 2009-2022, Panagiotis Christopoulos Charitos and contributors.
// All rights reserved.
// Code licensed under the BSD License.
// http://www.anki3d.org/LICENSE

// The 1st reduction is average depth. The next reductions are max depth

#pragma anki mutator WAVE_OPERATIONS 0 1

#pragma anki start comp
#include <AnKi/Shaders/Common.glsl>
#include <AnKi/Shaders/Include/MiscRendererTypes.h>

layout(local_size_x = 256) in;

layout(push_constant, std140) uniform b_pc
{
	DepthDownscaleUniforms u_unis;
};

layout(set = 0, binding = 0) uniform image2D u_dstImages[12u];
layout(set = 0, binding = 1) coherent uniform image2D u_dstImage5;

layout(set = 0, binding = 2) coherent buffer b_atomic
{
	U32 u_spdCounter;
};

layout(std430, set = 0, binding = 3) writeonly buffer b_cb
{
	F32 u_clientBuf[];
};

layout(set = 0, binding = 4) uniform sampler u_linearAnyClampSampler;
layout(set = 0, binding = 5) uniform texture2D u_srcTex;

// Include SPD
#define A_GPU 1
#define A_GLSL 1
#include <ThirdParty/FidelityFX/ffx_a.h>

shared AU1 s_spdCounter;
shared AF1 s_spdIntermediateR[16][16];

AF4 SpdLoadSourceImage(AU2 p, AU1 slice)
{
	const AF2 textureCoord = Vec2(p) * u_unis.m_srcTexSizeOverOne + u_unis.m_srcTexSizeOverOne;
	return AF4(textureLod(u_srcTex, u_linearAnyClampSampler, textureCoord, 0.0).r, 0.0, 0.0, 0.0);
}

AF4 SpdLoad(AU2 p, AU1 slice)
{
	return AF4(imageLoad(u_dstImage5, IVec2(p)).r, 0.0, 0.0, 0.0);
}

void SpdStore(AU2 p, AF4 value, AU1 mip, AU1 slice)
{
	if(mip == 5u)
	{
		imageStore(u_dstImage5, IVec2(p), Vec4(value.x, 0.0, 0.0, 0.0));
	}
	else
	{
		imageStore(u_dstImages[mip], IVec2(p), Vec4(value.x, 0.0, 0.0, 0.0));
	}

	// Store the last mip to the buffer as well
	if(mip == u_unis.m_mipmapCount - 1u)
	{
		const U32 idx = p.y * u_unis.m_lastMipWidth + p.x;
		u_clientBuf[idx] = value.x;
	}
}

void SpdIncreaseAtomicCounter(AU1 slice)
{
	s_spdCounter = atomicAdd(u_spdCounter, 1u);
}

AU1 SpdGetAtomicCounter()
{
	return s_spdCounter;
}

void SpdResetAtomicCounter(AU1 slice)
{
	u_spdCounter = 0u;
}

AF4 SpdLoadIntermediate(AU1 x, AU1 y)
{
	return AF4(s_spdIntermediateR[x][y], 0.0, 0.0, 0.0);
}

void SpdStoreIntermediate(AU1 x, AU1 y, AF4 value)
{
	s_spdIntermediateR[x][y] = value.x;
}

AF4 SpdReduce4(AF4 v0, AF4 v1, AF4 v2, AF4 v3)
{
	const F32 maxDepth = max(v0.x, max(v1.x, max(v2.x, v3.x)));
	return AF4(maxDepth, 0.0, 0.0, 0.0);
}

#define SPD_LINEAR_SAMPLER 1

#if WAVE_OPERATIONS == 0
#	define SPD_NO_WAVE_OPERATIONS 1
#endif

#include <ThirdParty/FidelityFX/ffx_spd.h>

void main()
{
	const U32 slice = 0u;
	const UVec2 offset = UVec2(0u);
	SpdDownsample(AU2(gl_WorkGroupID.xy), AU1(gl_LocalInvocationIndex), AU1(u_unis.m_mipmapCount),
				  AU1(u_unis.m_workgroupCount), slice, offset);
}

#pragma anki end