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

// if VARIANT==0 then the checkerboard pattern is (render on 'v'):
// -----
// |v| |
// | |v|
// -----

#pragma anki mutator VARIANT 0 1
#define EXTRA_REJECTION 0

#pragma anki start comp
#include <AnKi/Shaders/Functions.glsl>
#include <AnKi/Shaders/Pack.glsl>
#include <AnKi/Shaders/Include/SsrTypes.h>
#include <AnKi/Shaders/Tonemapping.glsl>
#include <AnKi/Shaders/SsRaymarching.glsl>

const UVec2 WORKGROUP_SIZE = UVec2(16, 16);
layout(local_size_x = WORKGROUP_SIZE.x, local_size_y = WORKGROUP_SIZE.y, local_size_z = 1) in;

layout(set = 0, binding = 0, rgba16f) uniform image2D out_img;

layout(set = 0, binding = 1, row_major) uniform u_
{
	SsrUniforms u_unis;
};

layout(set = 0, binding = 2) uniform sampler u_trilinearClampSampler;
layout(set = 0, binding = 3) uniform texture2D u_gbufferRt1;
layout(set = 0, binding = 4) uniform texture2D u_gbufferRt2;
layout(set = 0, binding = 5) uniform texture2D u_depthRt;
layout(set = 0, binding = 6) uniform texture2D u_lightBufferRt;

layout(set = 0, binding = 7) uniform sampler u_trilinearRepeatSampler;
layout(set = 0, binding = 8) uniform texture2D u_noiseTex;
const Vec2 NOISE_TEX_SIZE = Vec2(16.0);

void main()
{
	// Compute a global invocation ID that takes the checkerboard pattern into account
	IVec2 fixedGlobalInvocationId = IVec2(gl_GlobalInvocationID.xy);
	fixedGlobalInvocationId.x *= 2;
#if VARIANT == 0
	fixedGlobalInvocationId.x += ((fixedGlobalInvocationId.y + 1) & 1);
#else
	fixedGlobalInvocationId.x += ((fixedGlobalInvocationId.y + 0) & 1);
#endif

	if(fixedGlobalInvocationId.x >= I32(u_unis.m_framebufferSize.x)
	   || fixedGlobalInvocationId.y >= I32(u_unis.m_framebufferSize.y))
	{
		// Skip threads outside the writable image
		return;
	}

	const Vec2 uv = (Vec2(fixedGlobalInvocationId.xy) + 0.5) / Vec2(u_unis.m_framebufferSize);

	// Read part of the G-buffer
	const F32 roughness = readRoughnessFromGBuffer(u_gbufferRt1, u_trilinearClampSampler, uv);
	const Vec3 worldNormal = readNormalFromGBuffer(u_gbufferRt2, u_trilinearClampSampler, uv);

	// Get depth
	const F32 depth = textureLod(u_depthRt, u_trilinearClampSampler, uv, 0.0).r;

	// Get view pos
	const Vec4 viewPos4 = u_unis.m_invProjMat * Vec4(UV_TO_NDC(uv), depth, 1.0);
	const Vec3 viewPos = viewPos4.xyz / viewPos4.w;

	// Compute refl vector
	const Vec3 viewDir = normalize(viewPos);
	const Vec3 viewNormal = u_unis.m_normalMat * worldNormal;
	const Vec3 reflVec = reflect(viewDir, viewNormal);

	// Rand idx
	const Vec2 noiseUv = Vec2(u_unis.m_framebufferSize) / NOISE_TEX_SIZE * uv;
	const Vec2 noiseShift = 1.0 / NOISE_TEX_SIZE * F32(u_unis.m_frameCount % 4u);
	const F32 noise = textureLod(u_noiseTex, u_trilinearRepeatSampler, noiseUv + noiseShift, 0.0).r;

	// Do the heavy work
	Vec3 hitPoint;
	F32 hitAttenuation;
	const U32 lod = 0u;
	const U32 step = u_unis.m_firstStepPixels;
	const F32 stepf = F32(step);
	const F32 minStepf = stepf / 4.0;
	raymarchGroundTruth(viewPos, reflVec, uv, depth, u_unis.m_projMat, u_unis.m_maxSteps, u_depthRt,
						u_trilinearClampSampler, F32(lod), u_unis.m_depthBufferSize, step,
						U32((stepf - minStepf) * noise + minStepf), hitPoint, hitAttenuation);

#if EXTRA_REJECTION
	// Reject backfacing
	ANKI_BRANCH if(hitAttenuation > 0.0)
	{
		const Vec3 hitNormal =
			u_unis.m_normalMat * readNormalFromGBuffer(u_gbufferRt2, u_trilinearClampSampler, hitPoint.xy);
		F32 backFaceAttenuation;
		rejectBackFaces(reflVec, hitNormal, backFaceAttenuation);

		hitAttenuation *= backFaceAttenuation;
	}

	// Reject far from hit point
	ANKI_BRANCH if(hitAttenuation > 0.0)
	{
		const F32 depth = textureLod(u_depthRt, u_trilinearClampSampler, hitPoint.xy, 0.0).r;
		Vec4 viewPos4 = u_unis.m_invProjMat * Vec4(UV_TO_NDC(hitPoint.xy), depth, 1.0);
		const F32 actualZ = viewPos4.z / viewPos4.w;

		viewPos4 = u_unis.m_invProjMat * Vec4(UV_TO_NDC(hitPoint.xy), hitPoint.z, 1.0);
		const F32 hitZ = viewPos4.z / viewPos4.w;

		const F32 rejectionMeters = 1.0;
		const F32 diff = abs(actualZ - hitZ);
		const F32 distAttenuation = (diff < rejectionMeters) ? 1.0 : 0.0;
		hitAttenuation *= distAttenuation;
	}
#endif

	// Read the reflection
	Vec4 outColor;
	ANKI_BRANCH if(hitAttenuation > 0.0)
	{
		// Reproject the UV because you are reading the previous frame
		const Vec4 v4 = u_unis.m_prevViewProjMatMulInvViewProjMat * Vec4(UV_TO_NDC(hitPoint.xy), hitPoint.z, 1.0);
		hitPoint.xy = NDC_TO_UV(v4.xy / v4.w);

		// Compute the LOD based on the roughness
		const F32 lod = F32(u_unis.m_lightBufferMipCount - 1u) * roughness;

		// Read the light buffer
		outColor.rgb = textureLod(u_lightBufferRt, u_trilinearClampSampler, hitPoint.xy, lod).rgb;
		outColor.rgb = clamp(outColor.rgb, 0.0, FLT_MAX); // Fix the value just in case
		outColor.rgb *= hitAttenuation;
		outColor.a = 1.0 - hitAttenuation;
	}
	else
	{
		outColor = Vec4(0.0, 0.0, 0.0, 1.0);
	}

	// Store
	imageStore(out_img, fixedGlobalInvocationId, outColor);
}
#pragma anki end