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

#pragma once

#include "../RayTracing/RayTracingCapsuleLight.ush"
#include "../RayTracing/RayTracingSphereLight.ush"
#include "../RayTracing/RayTracingRectLight.ush"
#include "../RayTracing/RayTracingDirectionalLight.ush"

uint2 NumSamplesPerPixelDivideShift;

uint2 SampleCoordToDownsampledScreenCoord(uint2 SampleCoord)
{
	// SampleCoord / NumSamplesPerPixel
	return SampleCoord >> NumSamplesPerPixelDivideShift;
}

uint3 NumSamplesPerVoxelDivideShift;

uint3 SampleCoordToDownsampledVolumeCoord(uint3 SampleCoord)
{
	// SampleCoord / NumSamplesPerVoxel
	return SampleCoord >> NumSamplesPerVoxelDivideShift;
}

uint2 SampleCoordToScreenCoord(uint2 SampleCoord)
{
	uint2 DownsampledScreenCoord = SampleCoordToDownsampledScreenCoord(SampleCoord);
	return (DownsampledScreenCoord << DownsampleFactorMultShift) + GetSampleScreenCoordJitter(DownsampledScreenCoord);
}

bool GenerateShadowRay(
	FLightShaderParameters LightParameters,
	bool bSpotLight,
	bool bRectLight,
	bool bDirLight,
	float3 TranslatedWorldPosition,
	float3 WorldNormal,
	float2 RandSample,
	out float3 RayOrigin,
	out float3 RayDirection,
	out float RayTMin,
	out float RayTMax)
{
	float RayPdf = 0.0;
	if (bRectLight)
	{
		return GenerateRectLightOcclusionRay(
			LightParameters,
			TranslatedWorldPosition, WorldNormal,
			RandSample,
			/* out */ RayOrigin,
			/* out */ RayDirection,
			/* out */ RayTMin,
			/* out */ RayTMax,
			/* out */ RayPdf);
	}
	else if (bDirLight)
	{
		GenerateDirectionalLightOcclusionRay(
			LightParameters,
			TranslatedWorldPosition, WorldNormal,
			RandSample,
			/* out */ RayOrigin,
			/* out */ RayDirection,
			/* out */ RayTMin,
			/* out */ RayTMax);

		return true;
	}
	else
	{
		// point or spot light
		if (LightParameters.SourceLength > 0.0)
		{
			return GenerateCapsuleLightOcclusionRayWithSolidAngleSampling(
				LightParameters,
				TranslatedWorldPosition, WorldNormal,
				RandSample,
				/* out */ RayOrigin,
				/* out */ RayDirection,
				/* out */ RayTMin,
				/* out */ RayTMax,
				/* out */ RayPdf);
		}
		return GenerateSphereLightOcclusionRayWithSolidAngleSampling(
			LightParameters,
			TranslatedWorldPosition, WorldNormal,
			RandSample,
			/* out */ RayOrigin,
			/* out */ RayDirection,
			/* out */ RayTMin,
			/* out */ RayTMax,
			/* out */ RayPdf);
	}
}

struct FLightSampleTrace
{
	float3 Direction;
	float Distance;
};

/**
 * Compute shadow ray direction and distance for a given light sample.
 */
FLightSampleTrace GetLightSampleTrace(float3 TranslatedWorldPosition, uint LocalLightIndex, float2 LightSampleUV)
{
	FLightSampleTrace LightSampleTrace;
	LightSampleTrace.Direction = 0.0f;
	LightSampleTrace.Distance = 0.0f;

	if (LocalLightIndex != MAX_LOCAL_LIGHT_INDEX)
	{
		const FForwardLightData ForwardLightData = GetForwardLightData(LocalLightIndex, 0);
		const FDeferredLightData LightData = ConvertToDeferredLight(ForwardLightData);
		const FLightShaderParameters LightParameters = ConvertToLightShaderParameters(LightData);

		float3 Unused;
		float Unused2;

		bool bValid = GenerateShadowRay(
			LightParameters,
			LightData.bSpotLight,
			LightData.bRectLight,
			!LightData.bRadialLight,
			TranslatedWorldPosition,
			float3(0, 0, 0),
			LightSampleUV,
			Unused,
			LightSampleTrace.Direction,
			Unused2,
			LightSampleTrace.Distance);
		if (!bValid)
		{
			// MegaLight pipeline is not really set up to handle invalid samples.
			// Replace what we got with a valid sample so that the rest of the code can
			// still estimate a valid shadow fraction. Forcing the sample to be always
			// occluded would add noise in unshadowed regions, and forcing the sample to
			// count as un-occluded would add light leaks in occluded regions.
			// This seems like the simplest workaround for now.
			float3 ToLight = LightParameters.TranslatedWorldPosition - TranslatedWorldPosition;
			LightSampleTrace.Distance = length(ToLight);
			LightSampleTrace.Direction = ToLight * rcp(LightSampleTrace.Distance);
		}
	}

	return LightSampleTrace;
}

uint PackTraceTexel(uint2 SampleCoord)
{
	uint PackedTraceTexel = SampleCoord.y & 0xFFFF;
	PackedTraceTexel |= (SampleCoord.x & 0xFFFF) << 16;
	return PackedTraceTexel;
}

uint2 UnpackTraceTexel(uint PackedTraceTexel)
{
	uint2 SampleCoord;
	SampleCoord.x = PackedTraceTexel >> 16;
	SampleCoord.y = PackedTraceTexel & 0xFFFF;
	return SampleCoord;
}

uint PackTraceVoxel(uint3 SampleCoord)
{
	uint PackedTraceVoxel = SampleCoord.z & 0x3FF;
	PackedTraceVoxel |= (SampleCoord.y & 0x3FF) << 10;
	PackedTraceVoxel |= (SampleCoord.x & 0x3FF) << 20;
	return PackedTraceVoxel;
}

uint3 UnpackTraceVoxel(uint PackedTraceVoxel)
{
	uint3 SampleCoord;
	SampleCoord.x = (PackedTraceVoxel >> 20) & 0x3FF;
	SampleCoord.y = (PackedTraceVoxel >> 10) & 0x3FF;
	SampleCoord.z = PackedTraceVoxel & 0x3FF;
	return SampleCoord;
}