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

#pragma anki technique vert pixel

#include <AnKi/Shaders/QuadVert.hlsl>

#if ANKI_PIXEL_SHADER
#	include <AnKi/Shaders/PackFunctions.hlsl>
#	include <AnKi/Shaders/Functions.hlsl>
#	include <AnKi/Shaders/RtShadows.hlsl>
#	include <AnKi/Shaders/ClusteredShadingFunctions.hlsl>

ConstantBuffer<GlobalRendererConstants> g_globalConstants : register(b0);
StructuredBuffer<PointLight> g_pointLights : register(t0);
StructuredBuffer<SpotLight> g_spotLights : register(t1);
#	if INDIRECT_DIFFUSE_TEX
Texture2D<Vec4> g_indirectDiffuseTex : register(t2);
#	else
StructuredBuffer<GlobalIlluminationProbe> g_giProbes : register(t2);
#	endif
StructuredBuffer<Cluster> g_clusters : register(t4);

SamplerState g_nearestAnyClampSampler : register(s0);
SamplerState g_trilinearClampSampler : register(s1);

Texture2D<Vec4> g_gbuffer0Tex : register(t5);
Texture2D<Vec4> g_gbuffer1Tex : register(t6);
Texture2D<Vec4> g_gbuffer2Tex : register(t7);
Texture2D g_depthTex : register(t8);
Texture2D<Vec4> g_resolvedShadowsTex : register(t9);
Texture2D<Vec4> g_ssaoTex : register(t10);
Texture2D<Vec4> g_reflectionsTex : register(t11);
Texture2D<Vec4> g_integrationLut : register(t12);

// Common code for lighting
#	define LIGHTING_COMMON_BRDF() \
		const Vec3 frag2Light = light.m_position - worldPos; \
		const HVec3 l = normalize(frag2Light); \
		const HVec3 specC = specularIsotropicLobe(gbuffer.m_normal, gbuffer.m_f0, gbuffer.m_roughness, viewDir, l); \
		const HVec3 diffC = diffuseLobe(gbuffer.m_diffuse); \
		const F16 att = computeAttenuationFactor<F16>(light.m_radius, frag2Light); \
		F16 lambert = max(F16(0.0), dot(gbuffer.m_normal, l));

Vec4 main(VertOut input) : SV_TARGET0
{
	const Vec2 uv = input.m_uv;
	const Vec2 ndc = uvToNdc(uv);
	const UVec2 coord = input.m_svPosition;
	const F32 depth = g_depthTex[coord].r;

	if(depth == 1.0)
	{
		return 0.0;
	}

	// Get world position
	const Vec4 worldPos4 = mul(g_globalConstants.m_matrices.m_invertedViewProjectionJitter, Vec4(ndc, depth, 1.0));
	const Vec3 worldPos = worldPos4.xyz / worldPos4.w;

	const HVec3 viewDir = normalize(g_globalConstants.m_cameraPosition - worldPos);

	// Get the cluster
	Cluster cluster = getClusterFragCoord(g_clusters, g_globalConstants, Vec3(input.m_svPosition.xy, depth));

	// return clusterHeatmap(cluster, 1u << (U32)GpuSceneNonRenderableObjectType::kLight, 3);

	// Decode GBuffer
	GbufferInfo<F16> gbuffer = (GbufferInfo<F16>)0;
	unpackGBufferNoVelocity<F16>(g_gbuffer0Tex[coord], g_gbuffer1Tex[coord], g_gbuffer2Tex[coord], gbuffer);
	gbuffer.m_subsurface = max(gbuffer.m_subsurface, kSubsurfaceMin);

	// Apply SSAO
	const HVec4 ssaoAndBentNormals = g_ssaoTex.SampleLevel(g_trilinearClampSampler, uv, 0.0);
	const F16 ssao = ssaoAndBentNormals.w;
	const HVec3 bentNormal = ssaoAndBentNormals.xyz;
	gbuffer.m_diffuse *= ssao;

	// Ambient and emissive color
	HVec3 outColor = gbuffer.m_emission;

	// Indirect diffuse
#	if INDIRECT_DIFFUSE_TEX
	const HVec3 indirectCol = g_indirectDiffuseTex[coord];
	outColor += indirectCol * gbuffer.m_diffuse / kPi;
#	else
	const HVec3 probeColor = sampleGiProbes<F16>(cluster, g_giProbes, bentNormal, worldPos, g_trilinearClampSampler);
	outColor += probeColor * gbuffer.m_diffuse;
#	endif

	// Indirect specular
	{
		HVec3 refl = g_reflectionsTex[coord].xyz;

		// Apply the reflection
		const F16 NoV = max(0.0, dot(gbuffer.m_normal, viewDir));
		const Vec3 env = specularDFG<F16>(gbuffer.m_f0, gbuffer.m_roughness, g_integrationLut, g_trilinearClampSampler, NoV);
		refl *= env;

		outColor += refl;
	}

	// SM
	HVec4 resolvedSm = g_resolvedShadowsTex.SampleLevel(g_nearestAnyClampSampler, uv, 0.0);
	U32 resolvedSmIdx = 0u;

	// Dir light
	const DirectionalLight dirLight = g_globalConstants.m_directionalLight;
	if(dirLight.m_active)
	{
		F16 shadowFactor;
		if(dirLight.m_shadowCascadeCount)
		{
			shadowFactor = resolvedSm[0];
			++resolvedSmIdx;
		}
		else
		{
			shadowFactor = 1.0;
		}

		const HVec3 l = -dirLight.m_direction;

		const F16 lambert = max(gbuffer.m_subsurface, dot(l, gbuffer.m_normal));

		const HVec3 diffC = diffuseLobe(gbuffer.m_diffuse);
		const HVec3 specC = specularIsotropicLobe(gbuffer.m_normal, gbuffer.m_f0, gbuffer.m_roughness, viewDir, l);

		outColor += (diffC + specC) * dirLight.m_diffuseColor * (shadowFactor * lambert);
	}

	// Point lights
	U32 idx;
	[loop] while((idx = iteratePointLights(cluster)) != kMaxU32)
	{
		const PointLight light = g_pointLights[idx];

		LIGHTING_COMMON_BRDF();

		[branch] if(light.m_shadow)
		{
			const F16 shadow = resolvedSm[resolvedSmIdx++];
			lambert *= shadow;
		}

		outColor += (diffC + specC) * light.m_diffuseColor * (att * max(gbuffer.m_subsurface, lambert));
	}

	// Spot lights
	[loop] while((idx = iterateSpotLights(cluster)) != kMaxU32)
	{
		const SpotLight light = g_spotLights[idx];

		LIGHTING_COMMON_BRDF();

		const F16 spot = computeSpotFactor<F16>(l, light.m_outerCos, light.m_innerCos, light.m_direction);

		[branch] if(light.m_shadow)
		{
			const F16 shadow = resolvedSm[resolvedSmIdx++];
			lambert *= shadow;
		}

		outColor += (diffC + specC) * light.m_diffuseColor * (att * spot * max(gbuffer.m_subsurface, lambert));
	}

	if(any(isnan(outColor)) || any(isinf(outColor)))
	{
		outColor = 0.0;
	}

	return Vec4(outColor, 0.0);
}
#endif // ANKI_PIXEL_SHADER