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

#include "../Common.ush"
#include "../DeferredShadingCommon.ush"
#include "../SHCommon.ush"
#include "../ShaderPrint.ush"
#include "LumenCardCommon.ush"
#include "LumenTracingCommon.ush"
#include "LumenSoftwareRayTracing.ush"
#define RADIANCE_CACHE_DEPTH_TEST_SPHERE_PARALLAX 1
#include "LumenRadianceCacheCommon.ush"
#include "LumenVisualizeTraces.ush"
#include "LumenVisualize.ush"
#include "../ShaderPrint.ush"

uint3 CullGridSize;
uint CardGridPixelSizeShift;
int VisualizeTileIndex;
float VisualizeStepFactor;
float MinTraceDistance;
float MaxTraceDistance;
float MaxMeshSDFTraceDistanceForVoxelTracing;
float MaxMeshSDFTraceDistance;
float CardInterpolateInfluenceRadius;

RWTexture2D<float4> RWSceneColor;
Buffer<uint> CulledCardGridHeader;

#ifndef THREADGROUP_SIZE
#define THREADGROUP_SIZE 0
#endif

#ifdef VisualizeQuadsCS
[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, 1)]
void VisualizeQuadsCS(
	uint3 GroupId : SV_GroupID,
	uint3 DispatchThreadId : SV_DispatchThreadID,
	uint3 GroupThreadId : SV_GroupThreadID)
{
	uint2 TraceCoord = DispatchThreadId.xy;
	uint2 OutputScreenCoord = TraceCoord + OutputViewOffset;
	float2 ViewportUV = (TraceCoord + 0.5f) / OutputViewSize;
	uint2 PixelPos = ViewportUV * InputViewSize + InputViewOffset;

	// Clip rendering outside of bounds
	if (any(TraceCoord >= OutputViewSize))
	{
		return;
	}

	float2 ScreenPosition = float2(2.0f, -2.0f) * ViewportUV + float2(-1.0f, 1.0f);
	// We are rendering after TAAU, remove jitter
	float3 ScreenVector = ScreenVectorFromScreenRect(float4(ScreenPosition + View.TemporalAAJitter.xy, 1, 0));
	float3 CameraVector = normalize(ScreenVector);
	float Noise = InterleavedGradientNoise(PixelPos, View.StateFrameIndexMod8);

	float ConeAngle = View.EyeToPixelSpreadAngle;
	float TraceDistance = MaxTraceDistance;
	bool bCoveredByRadianceCache = false;

#if RADIANCE_CACHE
	FRadianceCacheCoverage Coverage = GetRadianceCacheCoverage(DFHackToFloat(PrimaryView.WorldCameraOrigin), CameraVector, .5f);
				
	if (Coverage.bValid)
	{
		TraceDistance = min(TraceDistance, Coverage.MinTraceDistanceBeforeInterpolation);
	}
#endif

	FConeTraceInput TraceInput;
	TraceInput.Setup(DFHackToFloat(PrimaryView.WorldCameraOrigin), PrimaryView.TranslatedWorldCameraOrigin, CameraVector, ConeAngle, .01f, MinTraceDistance, TraceDistance, VisualizeStepFactor);
	TraceInput.SDFStepFactor = lerp(.8f, 1.0f, Noise);
	TraceInput.bDitheredTransparency = false;
	TraceInput.DitherScreenCoord = PixelPos;
	TraceInput.bHiResSurface = VisualizeHiResSurface != 0 ? true : false;

	float3 Debug = -1;

	bool bContinueCardTracing = false;
#if TRACE_MESH_SDF || SCENE_TRACE_HEIGHTFIELDS
	bContinueCardTracing = true;
#endif

	TraceInput.VoxelTraceStartDistance = CalculateVoxelTraceStartDistance(MinTraceDistance, TraceDistance, MaxMeshSDFTraceDistance, bContinueCardTracing);

#if RADIANCE_CACHE
	TraceInput.VoxelTraceStartDistance = min(TraceInput.VoxelTraceStartDistance, TraceDistance);
#endif

	// The visualization pass is after TAAU, correct for being after the upscaler
	int2 RenderingViewportPixelPos = (PixelPos - InputViewOffset) * (int2)View.ViewSizeAndInvSize / InputViewSize;
	int2 CellCoord = clamp((int2)RenderingViewportPixelPos >> CardGridPixelSizeShift, int2(0, 0), (int2)CullGridSize.xy - 1);
	uint CardGridCellIndex = CellCoord.x + CellCoord.y * CullGridSize.x;

	TraceInput.NumMeshSDFs = NumGridCulledMeshSDFObjects[CardGridCellIndex];
	TraceInput.MeshSDFStartOffset = GridCulledMeshSDFObjectStartOffsetArray[CardGridCellIndex];
	TraceInput.CardInterpolateInfluenceRadius = CardInterpolateInfluenceRadius;
	//Debug = TraceInput.NumMeshSDFs / 20.0f;
	TraceInput.bCalculateHitVelocity = true;
	// Only expand the SDF surface when the ray traversal is far enough away to not self-intersect
	TraceInput.bExpandSurfaceUsingRayTimeInsteadOfMaxDistance = false;
	TraceInput.InitialMaxDistance = 0;

	FConeTraceResult TraceResult = (FConeTraceResult)0;
	TraceResult.Transparency = 1;
	TraceResult.OpaqueHitDistance = TraceInput.MaxTraceDistance;

#if TRACE_MESH_SDF
	ConeTraceLumenSceneCards(TraceInput, TraceResult);
#endif

#if SCENE_TRACE_HEIGHTFIELDS
	TraceInput.NumHeightfields = NumGridCulledHeightfieldObjects[CardGridCellIndex];
	TraceInput.HeightfieldStartOffset = GridCulledHeightfieldObjectStartOffsetArray[CardGridCellIndex];
	ConeTraceLumenSceneHeightfields(TraceInput, TraceResult);
#endif

#if TRACE_GLOBAL_SDF
	ConeTraceLumenSceneVoxels(TraceInput, TraceResult);
#endif

#if RADIANCE_CACHE
	TraceResult.Lighting += SampleRadianceCacheInterpolated(Coverage, DFHackToFloat(PrimaryView.WorldCameraOrigin), CameraVector, ConeAngle, /*RandomScalarForStochasticInterpolation*/ 0.5f).Radiance * TraceResult.Transparency;
	TraceResult.Transparency = 0.0f;
#else
	ApplySkylightToTraceResult(CameraVector, TraceResult);
#endif

	//TraceResult.Lighting += GetSkylightLeaking(CameraVector, TraceResult.OpaqueHitDistance);
	TraceResult.Lighting *= View.PreExposure;

	if (SampleHeightFog > 0)
	{
		float CoverageForFog = 1.0f;
		float3 OriginToCollider = CameraVector * TraceResult.OpaqueHitDistance;
		TraceResult.Lighting.rgb = GetFogOnLuminance(TraceResult.Lighting.rgb, CoverageForFog, TraceInput.ConeTranslatedOrigin, CameraVector, TraceResult.OpaqueHitDistance);
	}

	float4 FinalColor = 0.0f;

	#define TRACE_VISUALIZE_MODE 0
	#if TRACE_VISUALIZE_MODE == 0
		FinalColor = float4(TraceResult.Lighting + float3(0 * (saturate(TraceResult.NumOverlaps / 20.0f)), 0, 0), 0);
	#elif TRACE_VISUALIZE_MODE == 1
		FinalColor = float4(frac(TraceResult.OpaqueHitDistance.xxx / 100.0f), 0);
	#elif TRACE_VISUALIZE_MODE == 2
		FinalColor = float4(TraceResult.Lighting + TraceResult.Debug, 0);
	#elif TRACE_VISUALIZE_MODE == 3
		FinalColor = float4(abs(TraceResult.WorldVelocity / 100.0f), 0);
	#else 
		FinalColor = float4(TraceResult.NumSteps.xxx / 100.0f, 0);
	#endif

	if (any(Debug >= 0))
	{
		FinalColor = float4(Debug, 0);
	}

	LumenVisualizationFinalize(
		OutputScreenCoord,
		ViewportUV,
		CameraVector,
		FinalColor.xyz);

	RWSceneColor[OutputScreenCoord] = FinalColor;
}
#endif

#ifdef VisualizeCursorDataCS

StructuredBuffer<uint> DebugData;

[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, 1)]
void VisualizeCursorDataCS(
	uint3 GroupId : SV_GroupID,
	uint3 DispatchThreadId : SV_DispatchThreadID,
	uint3 GroupThreadId : SV_GroupThreadID)
{


	const FLumenSceneDebugData Data = ReadDebugData(DebugData);
	if (Data.bValid)
	{

		const FLumenCardData Card = GetLumenCardData(Data.CardIndex);
		const FLumenCardPageData CardPage = GetLumenCardPageData(Data.CardPageIndex);

		const uint TilePreviewSize = 256;
		uint2 OutputViewSize2 = 0;
		if (Card.LocalExtent.y < Card.LocalExtent.x)
		{
			const float Ratio = Card.LocalExtent.y / Card.LocalExtent.x;
			OutputViewSize2 = uint2(TilePreviewSize, TilePreviewSize * Ratio);
		}
		else
		{
			const float Ratio = Card.LocalExtent.x / Card.LocalExtent.y;
			OutputViewSize2 = uint2(TilePreviewSize * Ratio, TilePreviewSize);
		}
		uint2 OutputViewOffset2 = uint2(0, InputViewSize.y - OutputViewSize2.y);
		uint2 TraceCoord = DispatchThreadId.xy;
		uint2 OutputScreenCoord = TraceCoord + OutputViewOffset2;
		float2 ViewportUV = (TraceCoord + 0.5f) / OutputViewSize2;
		uint2 PixelPos = ViewportUV * InputViewSize + InputViewOffset;

		// Clip rendering outside of bounds
		if (any(TraceCoord >= OutputViewSize2))
		{
			return;
		}

		FShaderPrintContext Ctx = InitShaderPrintContext(all(DispatchThreadId == 0), uint2(50, 550));

		const float2 LocalSamplePosition = (float2(ViewportUV.x, 1 - ViewportUV.y) * 2.f - 1.f)  * Card.LocalExtent.xy;
		const FLumenCardSample Sample = ComputeSurfaceCacheSample(Card, Data.CardIndex, LocalSamplePosition, 1.f /*SampleRadius*/, true/*bHiResSurface*/);

		const float2 CardUV = ViewportUV;
		const float2 AtlasUV = Sample.PhysicalAtlasUV;

		const bool bInPage = all(AtlasUV >= CardPage.PhysicalAtlasUVRect.xy) && all(AtlasUV <= CardPage.PhysicalAtlasUVRect.zw);
		      
		const FLumenSurfaceCacheData SurfaceCacheData = GetSurfaceCacheData(Card, CardUV, AtlasUV);
		float3 OutColor = SurfaceCacheData.Albedo;

		Print(Ctx, TEXT("Card visualization"), FontEmerald); Newline(Ctx);
		const bool bAlbedo = AddCheckbox(Ctx, TEXT("Albedo"), false, FontWhite); Newline(Ctx);
		if (bAlbedo) { OutColor = Texture2DSampleLevel(AlbedoAtlas, View.SharedBilinearClampedSampler, AtlasUV, 0).xyz; }

		const bool bNormal = AddCheckbox(Ctx, TEXT("Normal"), false, FontWhite); Newline(Ctx);
		if (bNormal) { OutColor = (Texture2DSampleLevel(NormalAtlas, View.SharedBilinearClampedSampler, AtlasUV, 0).xyz + 1.f) * 0.5f; }

		const bool bEmissive = AddCheckbox(Ctx, TEXT("Emissive"), false, FontWhite); Newline(Ctx);
		if (bEmissive) { OutColor = Texture2DSampleLevel(EmissiveAtlas, View.SharedBilinearClampedSampler, AtlasUV, 0).xyz; }

		const bool bDirect = AddCheckbox(Ctx, TEXT("Direct lighting"), true, FontWhite); Newline(Ctx);
		if (bDirect) { OutColor = Texture2DSampleLevel(DirectLightingAtlas, View.SharedBilinearClampedSampler, AtlasUV, 0).xyz; }

		const bool bIndirect = AddCheckbox(Ctx, TEXT("Indirect lighting"), false, FontWhite); Newline(Ctx);
		if (bIndirect) { OutColor = Texture2DSampleLevel(IndirectLightingAtlas, View.SharedBilinearClampedSampler, AtlasUV, 0).xyz; }

		//const bool bFinal = AddCheckbox(Ctx, TEXT("Final lighting"), false, FontWhite); Newline(Ctx);
		//if (bFinal) { OutColor = Texture2DSampleLevel(FinalLightingAtlas, View.SharedBilinearClampedSampler, AtlasUV, 0).xyz; }

		const float Opacity = Texture2DSampleLevel(OpacityAtlas, View.SharedBilinearClampedSampler, AtlasUV, 0).x;
		OutColor *= Opacity;

		RWSceneColor[OutputScreenCoord] = float4(OutColor,1);

		AddQuadSS(Ctx, OutputViewOffset2, OutputViewOffset2 + OutputViewSize2, ColorYellow);
	}
}
#endif // VisualizeCursorDataCS

Buffer<float4> VisualizeTracesData;

struct FVisualizeTracesVertexOutput
{
	float3 TraceLighting : TEXCOORD0;	
};

void VisualizeTracesVS(
	uint VertexIndex : SV_VertexID,
	out FVisualizeTracesVertexOutput Output,
	out float4 OutPosition : SV_POSITION
	)
{
	uint TraceIndex = VertexIndex / 2;
	float3 WorldPosition = VisualizeTracesData[TraceIndex * VISUALIZE_TRACE_DATA_STRIDE + 0].xyz;

	if (VertexIndex & 1)
	{
		WorldPosition += VisualizeTracesData[TraceIndex * VISUALIZE_TRACE_DATA_STRIDE + 1].xyz;
	}

	OutPosition = mul(float4(WorldPosition, 1), DFHackToFloat(PrimaryView.WorldToClip));
	Output.TraceLighting = VisualizeTracesData[TraceIndex * VISUALIZE_TRACE_DATA_STRIDE + 2].xyz;
} 

void VisualizeTracesPS(
	FVisualizeTracesVertexOutput Input,
	in float4 SVPos : SV_POSITION,
	out float4 OutColor : SV_Target0)
{
	OutColor = float4(Input.TraceLighting, 0);
}

Texture2D SceneColorTexture;
Texture2D InputTexture;
float2 ToIntMulAdd;
float2 PreOutputMulAdd;
uint2 BitMaskShift;

void VisualizeBitFieldFloatTexturePS(
	in noperspective float4 UVAndScreenPos : TEXCOORD0,
	out float4 OutColor : SV_Target0)
{
	float2 SceneColorUV = UVAndScreenPos.zw * float2(0.5f, -0.5f) + float2(0.5f, 0.5f);
	float3 SceneColorRGB = Texture2DSampleLevel(SceneColorTexture, GlobalPointClampedSampler, SceneColorUV, 0).xyz;
	float SceneColorGrey = dot(SceneColorRGB, 1.0f) / 3.0f;
	
	float FloatValue = Texture2DSampleLevel(InputTexture, GlobalBilinearClampedSampler, UVAndScreenPos.xy, 0).x;
	uint IntValue = uint(FloatValue * ToIntMulAdd.x + ToIntMulAdd.y);
	uint BitField = (IntValue & BitMaskShift.x) >> BitMaskShift.y;
	float Amount = saturate(BitField * PreOutputMulAdd.x + PreOutputMulAdd.y);

	OutColor = float4(lerp(SceneColorGrey.xxx, float3(Amount, 0, 0), 0.8f), 0);
}

#ifdef VisualizeTracesCS
uint NumTracesToVisualize;

[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, 1)]
void VisualizeTracesCS(
	uint3 GroupId : SV_GroupID,
	uint3 GroupThreadId : SV_GroupThreadID,
	uint3 DispatchThreadId : SV_DispatchThreadID)
{
	uint TraceIndex = DispatchThreadId.x;

	if (TraceIndex < NumTracesToVisualize)
	{
		FShaderPrintContext Context = InitShaderPrintContext(true, float2(0.05, 0.05));

		float3 WorldPosition0 = VisualizeTracesData[TraceIndex * VISUALIZE_TRACE_DATA_STRIDE + 0].xyz;
		float3 WorldPosition1 = WorldPosition0 + VisualizeTracesData[TraceIndex * VISUALIZE_TRACE_DATA_STRIDE + 1].xyz;
		float3 TraceLighting = VisualizeTracesData[TraceIndex * VISUALIZE_TRACE_DATA_STRIDE + 2].xyz;

		float3 LineColor = VisualizeTonemap(TraceLighting);

		AddLineWS(Context, WorldPosition0, WorldPosition1, float4(LineColor, 1.0f));
	}
}
#endif