// Copyright Epic Games, Inc. All Rights Reserved.

#include "../Common.ush"
#include "../SHCommon.ush"
#include "LumenTracingCommon.ush"
#include "LumenRadianceCacheCommon.ush"

#define VISUALIZE_MODE_RADIANCE			1
#define VISUALIZE_MODE_SKY_VISIBILITY	2

float4 ClipmapCornerTWSAndCellSizeForVisualization;

float VisualizeProbeRadiusScale;
uint ProbeClipmapIndex;

struct FVisualizeRadianceCacheVSToPS
{
	nointerpolation float3 ProbeCoord : TEXCOORD0;
	nointerpolation float4 CellSphere : TEXCOORD1;
	float3 PositionTWS : TEXCOORD2;
};

float3 GetProbeTranslatedWorldPositionForVisualization(uint3 ProbeCoord, uint ProbeIndex)
{
	const float3 ClipmapCornerTWS = ClipmapCornerTWSAndCellSizeForVisualization.xyz;
	const float CellSize = ClipmapCornerTWSAndCellSizeForVisualization.w;

	const float3 CornerTranslatedWorldPosition = ClipmapCornerTWS;

	const float3 CornerToProbe = ((ProbeCoord + 0.5f) * CellSize);
	return CornerTranslatedWorldPosition + CornerToProbe + ProbeWorldOffset[ProbeIndex].xyz;
}

void VisualizeRadianceCacheVS(
	uint VertexId : SV_VertexID,
	uint InstanceId : SV_InstanceID,
	out FVisualizeRadianceCacheVSToPS Output,
	out float4 OutPosition : SV_POSITION
)
{
	float VisualizeRadius = VisualizeProbeRadiusScale * GetRadianceProbeClipmapCellSize(ProbeClipmapIndex);

	float3 LocalPosition;
	LocalPosition.x = VertexId & 0x1 ? 1.0f : -1.0f;
	LocalPosition.y = VertexId & 0x2 ? 1.0f : -1.0f;
	LocalPosition.z = VertexId & 0x4 ? 1.0f : -1.0f;
	LocalPosition *= VisualizeRadius;

	uint3 ProbeCoord;
	ProbeCoord.x = InstanceId % RadianceProbeClipmapResolution;
	ProbeCoord.y = (InstanceId / RadianceProbeClipmapResolution) % RadianceProbeClipmapResolution;
	ProbeCoord.z = InstanceId / (RadianceProbeClipmapResolution * RadianceProbeClipmapResolution);

	uint ProbeIndex = GetProbeIndexFromIndirectionTexture(ProbeCoord, ProbeClipmapIndex);

	float3 ProbeTranslatedWorldCenter = GetProbeTranslatedWorldPositionForVisualization(ProbeCoord, ProbeIndex);

	Output.CellSphere.xyz = ProbeTranslatedWorldCenter;
	Output.CellSphere.w = VisualizeRadius;
	Output.PositionTWS = LocalPosition + ProbeTranslatedWorldCenter;
	Output.ProbeCoord = ProbeCoord;
	OutPosition = mul(float4(Output.PositionTWS, 1.0f), PrimaryView.TranslatedWorldToClip);

	if (ProbeIndex == INVALID_PROBE_INDEX)
	{
		OutPosition = 0;
	}
}

void VisualizeRadianceCachePS(
	FVisualizeRadianceCacheVSToPS Input,
	out float4 OutColor : SV_Target0
)
{
	float3 RayDirection = Input.PositionTWS;
	float2 SphereIntersections = RayIntersectSphere(0, RayDirection, Input.CellSphere);
	float3 IntersectionPosition = SphereIntersections.x * RayDirection;
	clip(SphereIntersections.x);

	float3 SphereNormal = normalize(IntersectionPosition - Input.CellSphere.xyz);

	uint ProbeIndex = GetProbeIndexFromIndirectionTexture(Input.ProbeCoord, ProbeClipmapIndex);
	float3 Lighting = 0.0f;

	#if VISUALIZE_MODE == VISUALIZE_MODE_SKY_VISIBILITY
	{
		FRadianceCacheSample RadianceCacheSample = SampleRadianceCacheProbe(ProbeIndex, SphereNormal, 0);
		Lighting = RadianceCacheSample.SkyVisibility;
	}
	#else // VISUALIZE_MODE == VISUALIZE_MODE_RADIANCE
	{
		#if RADIANCE_CACHE_IRRADIANCE
		{
			float Albedo = .3f;
			Lighting = SampleIrradianceCacheProbe(ProbeIndex, SphereNormal) * Albedo;

			#define DEBUG_VISUALIZE_PROBES_WITH_OFFSETS 0
			#if DEBUG_VISUALIZE_PROBES_WITH_OFFSETS
				if (ProbeWorldOffset[ProbeIndex].w > 0)
				{
					Lighting.g += .1f;
				}
			#endif

			#define DEBUG_VISUALIZE_PROBE_OCCLUSION 0
			#if DEBUG_VISUALIZE_PROBE_OCCLUSION
				Lighting = float3(SampleProbeOcclusion(ProbeIndex, SphereNormal).xxx / 1000.0f);
			#endif
		}
		#else
		{
			FRadianceCacheSample RadianceCacheSample = SampleRadianceCacheProbe(ProbeIndex, SphereNormal, 0);
			RadianceCacheSample = RadianceCacheResolveSky(RadianceCacheSample, SphereNormal);
			Lighting = RadianceCacheSample.Radiance;
		}
		#endif
	}
	#endif

	if (ProbeIndex == INVALID_PROBE_INDEX)
	{
		Lighting = float3(0, 1000.0f, 0);
	}

	Lighting *= View.PreExposure;

	OutColor = float4(Lighting, 1.0f);
}