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

// Calculates the motion vectors that will be used to sample from the previous frame

#pragma anki technique vert pixel comp

#include <AnKi/Shaders/QuadVert.hlsl>

#if ANKI_COMPUTE_SHADER || ANKI_PIXEL_SHADER
#	include <AnKi/Shaders/Functions.hlsl>

SamplerState g_nearesetAnyClampSampler : register(s0);
Texture2D g_currentDepthTex : register(t0);
Texture2D g_velocityTex : register(t1);

struct Constants
{
	Mat4 m_currentViewProjMat;
	Mat4 m_currentInvViewProjMat;
	Mat4 m_prevViewProjMat;
};

ConstantBuffer<Constants> g_consts : register(b0);

#	if ANKI_COMPUTE_SHADER
RWTexture2D<Vec2> g_motionVectorsStorageTex : register(u0);
#	endif

#	if ANKI_COMPUTE_SHADER
[numthreads(8, 8, 1)] void main(UVec2 svDispatchThreadId : SV_DISPATCHTHREADID)
#	else
struct PixelOut
{
	Vec2 m_motionVectors : SV_TARGET0;
};

PixelOut main(VertOut input)
#	endif
{
#	if ANKI_COMPUTE_SHADER
	Vec2 texSize;
	g_motionVectorsStorageTex.GetDimensions(texSize.x, texSize.y);
	const Vec2 uv = (Vec2(svDispatchThreadId) + 0.5f) / texSize;
#	else
	const Vec2 uv = input.m_uv;
#	endif

	const Vec2 velocity = g_velocityTex.SampleLevel(g_nearesetAnyClampSampler, uv, 0.0).rg;

	Vec2 historyUv;
	if(velocity.x != 1.0)
	{
		historyUv = uv + velocity;
	}
	else
	{
		// Don't use a reprojection matrix or other kind of optimizations. Due to numerical precision it produces slightly off result. Big enough to
		// create slight visual issues. Do it the hard way.
		const F32 depth = g_currentDepthTex.SampleLevel(g_nearesetAnyClampSampler, uv, 0.0).r;
		const Vec4 v4 = mul(g_consts.m_currentInvViewProjMat, Vec4(uvToNdc(uv), depth, 1.0));
		const Vec3 worldPos = v4.xyz / v4.w;

		Vec4 clipPos = mul(g_consts.m_currentViewProjMat, Vec4(worldPos, 1.0));
		clipPos.xy /= clipPos.w;

		Vec4 prevClipPos = mul(g_consts.m_prevViewProjMat, Vec4(worldPos, 1.0));
		prevClipPos.xy /= prevClipPos.w;

		const Vec2 diff = ndcToUv(prevClipPos.xy) - ndcToUv(clipPos.xy);
		historyUv = uv + diff;
	}

	// Write out
#	if ANKI_COMPUTE_SHADER
	g_motionVectorsStorageTex[svDispatchThreadId] = historyUv - uv;
#	else
	PixelOut output;
	output.m_motionVectors = historyUv - uv;
	return output;
#	endif
}
#endif // ANKI_COMPUTE_SHADER || ANKI_PIXEL_SHADER