// 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 technique RtMaterialFetch rgen

#include <AnKi/Shaders/RtMaterialFetch.hlsl>
#include <AnKi/Shaders/Include/GpuSceneTypes.h>
#include <AnKi/Shaders/PackFunctions.hlsl>
#include <AnKi/Shaders/ImportanceSampling.hlsl>
#include <AnKi/Shaders/Functions.hlsl>
#include <AnKi/Shaders/LightFunctions.hlsl>

// Config and consts
constexpr Bool kTryShadowmapFirst = true;
constexpr F32 kTMax = 1000.0;

// Functions
Vec3 getDiffuseIndirect(StructuredBuffer<GpuSceneGlobalIlluminationProbe> giProbes, Vec3 worldPos, Vec3 worldNormal,
						SamplerState linearAnyClampSampler)
{
	const U32 probeCount = getStructuredBufferElementCount(giProbes);
	U32 i;
	for(i = 0; i < probeCount; ++i)
	{
		if(any(worldPos >= giProbes[i].m_aabbMax) || any(worldPos <= giProbes[i].m_aabbMin))
		{
			continue;
		}
		else
		{
			break;
		}
	}

	const Bool probeFound = (i != probeCount);
	if(probeFound)
	{
		const GpuSceneGlobalIlluminationProbe probe = giProbes[i];
		return sampleGlobalIllumination<F32>(worldPos, worldNormal, probe, getBindlessTexture3DVec4(probe.m_volumeTexture), linearAnyClampSampler);
	}
	else
	{
		return 0.0;
	}
}

Vec3 lightShading(Vec3 rayOrigin, Vec3 rayDir, Vec3 hitPos, Vec3 hitNormal, Vec3 emission, Vec3 diffuse, Bool isSky)
{
	Vec3 color = 0;

	if(isSky)
	{
		color = sampleSkyCheap<F32>(g_globalRendererConstants.m_sky, rayDir, g_linearAnyClampSampler);
	}
	else
	{
		const DirectionalLight dirLight = g_globalRendererConstants.m_directionalLight;

		// Trace shadow
		Vec4 vv4 = mul(g_globalRendererConstants.m_matrices.m_viewProjection, Vec4(hitPos, 1.0));
		vv4.xy /= vv4.w;
		const Bool bInsideFrustum = all(vv4.xy > -1.0) && all(vv4.xy < 1.0) && vv4.w > 0.0;

		F32 shadow;
		if(bInsideFrustum && kTryShadowmapFirst)
		{
			const F32 negativeZViewSpace = -mul(g_globalRendererConstants.m_matrices.m_view, Vec4(hitPos, 1.0)).z;
			const U32 shadowCascadeCount = dirLight.m_shadowCascadeCount;

			const U32 cascadeIdx = computeShadowCascadeIndex(negativeZViewSpace, dirLight.m_shadowCascadeDistances, shadowCascadeCount);

			shadow = computeShadowFactorDirLight<F32>(dirLight, cascadeIdx, hitPos, g_shadowAtlasTex, g_shadowSampler);
		}
		else
		{
			constexpr U32 qFlags = RAY_FLAG_FORCE_OPAQUE | RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES | RAY_FLAG_ACCEPT_FIRST_HIT_AND_END_SEARCH;
			RayQuery<qFlags> q;
			const U32 flags = RAY_FLAG_FORCE_OPAQUE;
			const U32 cullMask = 0xFFu;
			RayDesc ray;
			ray.Origin = hitPos;
			ray.TMin = 0.01;
			ray.Direction = -dirLight.m_direction;
			ray.TMax = kTMax;
			q.TraceRayInline(g_tlas, qFlags, cullMask, ray);
			q.Proceed();
			shadow = (q.CommittedStatus() == COMMITTED_TRIANGLE_HIT) ? 0.0 : 1.0;
		}

		// Do simple light shading
		color = emission;

		const Vec3 indirectDiffuse = getDiffuseIndirect(g_giProbes, hitPos, hitNormal, g_linearAnyClampSampler);
		// color += diffuse * indirectDiffuse;

		const Vec3 l = -dirLight.m_direction;
		const F32 lambert = max(0.0, dot(l, hitNormal));
		const Vec3 diffC = diffuseLobe(diffuse);
		color += diffC * dirLight.m_diffuseColor * lambert * shadow;
	}

	return color;
}

[Shader("raygeneration")] void main()
{
	Vec2 outSize;
	g_colorAndPdfTex.GetDimensions(outSize.x, outSize.y);

	const UVec2 coord = DispatchRaysIndex().xy;
	const Vec2 uv = Vec2(coord) / outSize;

	const F32 depth = g_depthTex[coord].x;
	const Vec4 rt2 = g_gbufferRt2[coord];
	const Vec3 worldNormal = unpackNormalFromGBuffer(rt2);

	const Vec4 v4 = mul(g_globalRendererConstants.m_matrices.m_invertedViewProjectionJitter, Vec4(uvToNdc(uv), depth, 1.0));
	const Vec3 worldPos = v4.xyz / v4.w;

	// Rand
	const UVec3 seed = rand3DPCG16(UVec3(coord, g_globalRendererConstants.m_frame % 8u));
	const Vec2 randFactors = hammersleyRandom16(g_globalRendererConstants.m_frame % 64u, 64u, seed);

	Vec3 rayOrigin = worldPos;
	Vec3 normal = worldNormal;

	Vec3 outColor = 0.0;
	Vec3 diffuse = 1.0;
	[unroll] for(U32 bounce = 0; bounce < 3; ++bounce)
	{
		const Mat3 tbn = rotationFromDirection(normal);
		const Vec3 rayDir = normalize(mul(tbn, hemisphereSampleCos(randFactors)));

		RtMaterialFetchRayPayload payload;
		payload.m_textureLod = 100.0;
		const U32 flags = RAY_FLAG_FORCE_OPAQUE;
		const U32 sbtRecordOffset = 0u;
		const U32 sbtRecordStride = 0u;
		const U32 missIndex = 0u;
		const U32 cullMask = 0xFFu;
		RayDesc ray;
		ray.Origin = rayOrigin;
		ray.TMin = 0.01;
		ray.Direction = rayDir;
		ray.TMax = kTMax;
		TraceRay(g_tlas, flags, cullMask, sbtRecordOffset, sbtRecordStride, missIndex, ray, payload);

		F32 rayT = payload.m_rayT;
		const Bool hasHitSky = rayT < 0.0;
		if(hasHitSky)
		{
			rayT = kTMax;
		}

		const Vec3 hitPos = rayOrigin + rayDir * rayT;

		outColor += diffuse * lightShading(worldPos, rayDir, hitPos, payload.m_worldNormal, payload.m_emission, payload.m_diffuseColor, hasHitSky);

		if(hasHitSky)
		{
			break;
		}

		rayOrigin = hitPos;
		normal = payload.m_worldNormal;
		diffuse = payload.m_diffuseColor;
	}

	g_colorAndPdfTex[coord] = Vec4(outColor, 0.0);
	// g_colorAndPdfTex[coord] = Vec4(lerp(outColor, g_colorAndPdfTex[coord].xyz, 0.98), 0.0);
}