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

/*=============================================================================
	ReflectionEnvironmentComposite
=============================================================================*/

#pragma once

#include "LargeWorldCoordinates.ush"
#include "LightGridCommon.ush"

float3 CompositeReflectionCapturesAndSkylightTWS(
	float CompositeAlpha, 
	float3 TranslatedWorldPosition, 
	float3 RayDirection, 
	float Roughness, 
	float IndirectIrradiance, 
	float IndirectSpecularOcclusion,
	float3 ExtraIndirectSpecular,
	uint NumCapturesAffectingTile,
	uint CaptureDataStartIndex, 
	int SingleCaptureIndex,
	bool bCompositeSkylight)
{
	float Mip = ComputeReflectionCaptureMipFromRoughness(Roughness, View.ReflectionCubemapMaxMip);
	float4 ImageBasedReflections = float4(0, 0, 0, CompositeAlpha);
	float2 CompositedAverageBrightness = float2(0.0f, 1.0f);

#if REFLECTION_COMPOSITE_USE_BLENDED_REFLECTION_CAPTURES
	// Accumulate reflections from captures affecting this tile, applying largest captures first so that the smallest ones display on top
	LOOP
	for (uint TileCaptureIndex = 0; TileCaptureIndex < NumCapturesAffectingTile; TileCaptureIndex++) 
	{
		BRANCH
		if (ImageBasedReflections.a < 0.001)
		{
			break;
		}

		uint CaptureIndex = 0;
#ifdef REFLECTION_COMPOSITE_NO_CULLING_DATA
		CaptureIndex = TileCaptureIndex;	// Go from 0 to NumCapturesAffectingTile as absolute index in capture array
#else
		CaptureIndex = GetCulledLightDataGrid(CaptureDataStartIndex + TileCaptureIndex);
#endif

		FDFVector3 CaptureWorldPosition = MakeDFVector3(GetReflectionPositionAndRadius(CaptureIndex).xyz, GetReflectionPositionLow(CaptureIndex).xyz);
		float3 CaptureTranslatedWorldPosition = DFFastAddDemote(CaptureWorldPosition, ResolvedView.PreViewTranslation);
		float CaptureRadius = GetReflectionPositionAndRadius(CaptureIndex).w;

		float4 CaptureProperties = GetReflectionCaptureProperties(CaptureIndex);

		float3 CaptureVector = TranslatedWorldPosition - CaptureTranslatedWorldPosition;
		float CaptureVectorLength = sqrt(dot(CaptureVector, CaptureVector));		
		float NormalizedDistanceToCapture = saturate(CaptureVectorLength / CaptureRadius);

		BRANCH
		if (CaptureVectorLength < CaptureRadius)
		{
			float3 ProjectedCaptureVector = RayDirection;
			float4 CaptureOffsetAndAverageBrightness = GetReflectionCaptureOffsetAndAverageBrightness(CaptureIndex);

			// Fade out based on distance to capture
			float DistanceAlpha = 0;
			
			#define PROJECT_ONTO_SHAPE 1
			#if PROJECT_ONTO_SHAPE
				#if REFLECTION_COMPOSITE_HAS_BOX_CAPTURES
					#if REFLECTION_COMPOSITE_HAS_SPHERE_CAPTURES
					// Box
					BRANCH if (CaptureProperties.b > 0)
					#endif
					{
						ProjectedCaptureVector = GetLookupVectorForBoxCapture(RayDirection, TranslatedWorldPosition, float4(CaptureTranslatedWorldPosition, CaptureRadius),
																			 GetReflectionBoxTransform(CaptureIndex), GetReflectionBoxScales(CaptureIndex), CaptureOffsetAndAverageBrightness.xyz, DistanceAlpha);
					}
				#endif

				#if REFLECTION_COMPOSITE_HAS_SPHERE_CAPTURES
					// Sphere
					#if REFLECTION_COMPOSITE_HAS_BOX_CAPTURES
					else
					#endif
					{
						ProjectedCaptureVector = GetLookupVectorForSphereCapture(RayDirection, TranslatedWorldPosition, float4(CaptureTranslatedWorldPosition, CaptureRadius), NormalizedDistanceToCapture, CaptureOffsetAndAverageBrightness.xyz, DistanceAlpha);
					}
				#endif
			#else 
				DistanceAlpha = 1.0;
			#endif //PROJECT_ONTO_SHAPE

			float CaptureArrayIndex = CaptureProperties.g;

			{
				float4 Sample = ReflectionStruct.ReflectionCubemap.SampleLevel(ReflectionStruct.ReflectionCubemapSampler, float4(ProjectedCaptureVector, CaptureArrayIndex), Mip);

				Sample.rgb *= CaptureProperties.r;
				Sample *= DistanceAlpha;

				// Under operator (back to front)
				ImageBasedReflections.rgb += Sample.rgb * ImageBasedReflections.a * IndirectSpecularOcclusion;
				ImageBasedReflections.a *= 1 - Sample.a;

				float AverageBrightness = CaptureOffsetAndAverageBrightness.w;
				CompositedAverageBrightness.x += AverageBrightness * DistanceAlpha * CompositedAverageBrightness.y;
				CompositedAverageBrightness.y *= 1 - DistanceAlpha;
			}
		}
	}

#else

	float3 ProjectedCaptureVector = RayDirection;

	FDFVector3 SingleCaptureWorldPosition = MakeDFVector3(GetReflectionPositionAndRadius(SingleCaptureIndex).xyz, GetReflectionPositionLow(SingleCaptureIndex).xyz);
	float3 SingleCaptureTranslatedWorldPosition = DFFastAddDemote(SingleCaptureWorldPosition, ResolvedView.PreViewTranslation);
	float SingleCaptureRadius = GetReflectionPositionAndRadius(SingleCaptureIndex).w;

	float4 SingleCaptureOffsetAndAverageBrightness = GetReflectionCaptureOffsetAndAverageBrightness(SingleCaptureIndex);
	float SingleCaptureBrightness = GetReflectionCaptureProperties(SingleCaptureIndex).x;
	float SingleCaptureArrayIndex = GetReflectionCaptureProperties(SingleCaptureIndex).y;

	#define APPROXIMATE_CONTINUOUS_SINGLE_CAPTURE_PARALLAX 0
	#if APPROXIMATE_CONTINUOUS_SINGLE_CAPTURE_PARALLAX
		float3 CaptureVector = TranslatedWorldPosition - SingleCaptureTranslatedWorldPosition;
		float CaptureVectorLength = sqrt(dot(CaptureVector, CaptureVector));		
		float NormalizedDistanceToCapture = saturate(CaptureVectorLength / SingleCaptureRadius);

		float UnusedDistanceAlpha = 0;
		ProjectedCaptureVector = GetLookupVectorForSphereCapture(RayDirection, TranslatedWorldPosition, float4(SingleCaptureTranslatedWorldPosition, SingleCaptureRadius), NormalizedDistanceToCapture, SingleCaptureOffsetAndAverageBrightness.xyz, UnusedDistanceAlpha);
		
		float x = saturate(NormalizedDistanceToCapture);
		float DistanceAlpha = 1 - x * x * (3 - 2 * x);
		// Lerp between sphere parallax corrected and infinite based on distance to shape
		ProjectedCaptureVector = lerp(RayDirection, normalize(ProjectedCaptureVector), DistanceAlpha);
	#endif

	float4 Sample = TextureCubeArraySampleLevel(ReflectionStruct.ReflectionCubemap, ReflectionStruct.ReflectionCubemapSampler, ProjectedCaptureVector, SingleCaptureArrayIndex, Mip);

	Sample.rgb *= SingleCaptureBrightness;
	ImageBasedReflections = float4(Sample.rgb, 1 - Sample.a);

	float AverageBrightness = SingleCaptureOffsetAndAverageBrightness.w;
	CompositedAverageBrightness.x += AverageBrightness * CompositedAverageBrightness.y;
	CompositedAverageBrightness.y = 0;
#endif

	// Apply indirect lighting scale while we have only accumulated reflection captures
	ImageBasedReflections.rgb *= View.PrecomputedIndirectSpecularColorScale;
	CompositedAverageBrightness.x *= Luminance( View.PrecomputedIndirectSpecularColorScale );

#if ENABLE_SKY_LIGHT

	BRANCH  
	if (ReflectionStruct.SkyLightParameters.y > 0 && bCompositeSkylight)
	{
		float SkyAverageBrightness = 1.0f;

		#if REFLECTION_COMPOSITE_SUPPORT_SKYLIGHT_BLEND
			float3 SkyLighting = GetSkyLightReflectionSupportingBlend(RayDirection, Roughness, SkyAverageBrightness);
		#else
			float3 SkyLighting = GetSkyLightReflection(RayDirection, Roughness, SkyAverageBrightness);
		#endif

		// Normalize for static skylight types which mix with lightmaps, material ambient occlusion as well as diffuse/specular occlusion.
		bool bNormalize = ReflectionStruct.SkyLightParameters.z < 1 && ALLOW_STATIC_LIGHTING;

		FLATTEN
		if (bNormalize)
		{
			ImageBasedReflections.rgb += ImageBasedReflections.a * SkyLighting * IndirectSpecularOcclusion;
			CompositedAverageBrightness.x += SkyAverageBrightness * CompositedAverageBrightness.y;
		}
		else
		{
			ExtraIndirectSpecular += SkyLighting * IndirectSpecularOcclusion;
		}
	}
#endif

#if ALLOW_STATIC_LIGHTING
	ImageBasedReflections.rgb *= ComputeMixingWeight(IndirectIrradiance, CompositedAverageBrightness.x, Roughness);
#endif

	ImageBasedReflections.rgb += ImageBasedReflections.a * ExtraIndirectSpecular;

	return ImageBasedReflections.rgb;
}