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

/*=============================================================================
	DecalRenderingShared.cpp
=============================================================================*/

#include "DecalRenderingShared.h"
#include "StaticBoundShaderState.h"
#include "Components/DecalComponent.h"
#include "GlobalShader.h"
#include "MaterialShaderType.h"
#include "MaterialShader.h"
#include "DebugViewModeRendering.h"
#include "ScenePrivate.h"
#include "PipelineStateCache.h"

static TAutoConsoleVariable<float> CVarDecalFadeScreenSizeMultiplier(
	TEXT("r.Decal.FadeScreenSizeMult"),
	1.0f,
	TEXT("Control the per decal fade screen size. Multiplies with the per-decal screen size fade threshold.")
	TEXT("  Smaller means decals fade less aggressively.")
	);

FTransientDecalRenderData::FTransientDecalRenderData(const FScene& InScene, const FDeferredDecalProxy& InDecalProxy, float InConservativeRadius)
	: Proxy(InDecalProxy)
	, FadeAlpha(1.0f)
	, ConservativeRadius(InConservativeRadius)
{
	MaterialProxy = Proxy.DecalMaterial->GetRenderProxy();
	MaterialResource = &MaterialProxy->GetMaterialWithFallback(InScene.GetFeatureLevel(), MaterialProxy);
	check(MaterialProxy && MaterialResource);
	BlendDesc = DecalRendering::ComputeDecalBlendDesc(InScene.GetShaderPlatform(), *MaterialResource);
}

/**
 * A vertex shader for projecting a deferred decal onto the scene.
 */
class FDeferredDecalVS : public FGlobalShader
{
	DECLARE_GLOBAL_SHADER(FDeferredDecalVS);
	SHADER_USE_PARAMETER_STRUCT(FDeferredDecalVS, FGlobalShader);

	BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
		SHADER_PARAMETER(FMatrix44f, FrustumComponentToClip)
		SHADER_PARAMETER_STRUCT_REF(FPrimitiveUniformShaderParameters, PrimitiveUniformBuffer)
	END_SHADER_PARAMETER_STRUCT()
};

IMPLEMENT_GLOBAL_SHADER(FDeferredDecalVS, "/Engine/Private/DeferredDecal.usf", "MainVS" ,SF_Vertex);

/**
 * A pixel shader for projecting a deferred decal onto the scene.
 */
class FDeferredDecalPS : public FMaterialShader
{
	DECLARE_SHADER_TYPE(FDeferredDecalPS,Material);

public:
	static bool ShouldCompilePermutation(const FMaterialShaderPermutationParameters& Parameters)
	{
		return (Parameters.MaterialParameters.MaterialDomain == MD_DeferredDecal) &&
			DecalRendering::GetBaseRenderStage(DecalRendering::ComputeDecalBlendDesc(Parameters.Platform, Parameters.MaterialParameters)) != EDecalRenderStage::None;
	}

	static void ModifyCompilationEnvironment(const FMaterialShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment)
	{
		FMaterialShader::ModifyCompilationEnvironment(Parameters, OutEnvironment);
		DecalRendering::ModifyCompilationEnvironment(DecalRendering::ComputeDecalBlendDesc(Parameters.Platform, Parameters.MaterialParameters), EDecalRenderStage::None, OutEnvironment);
	}

	FDeferredDecalPS() {}
	FDeferredDecalPS(const ShaderMetaType::CompiledShaderInitializerType& Initializer)
		: FMaterialShader(Initializer)
	{
		DecalTilePosition.Bind(Initializer.ParameterMap, TEXT("DecalTilePosition"));
		SvPositionToDecal.Bind(Initializer.ParameterMap,TEXT("SvPositionToDecal"));
		DecalToWorld.Bind(Initializer.ParameterMap,TEXT("DecalToWorld"));
		DecalToWorldInvScale.Bind(Initializer.ParameterMap, TEXT("DecalToWorldInvScale"));
		DecalOrientation.Bind(Initializer.ParameterMap,TEXT("DecalOrientation"));
		DecalParams.Bind(Initializer.ParameterMap, TEXT("DecalParams"));
	}

	void SetParameters(FRHICommandList& RHICmdList, const FViewInfo& View, const FMaterialRenderProxy* MaterialProxy, const FDeferredDecalProxy& DecalProxy, const float FadeAlphaValue=1.0f)
	{
		FRHIPixelShader* ShaderRHI = RHICmdList.GetBoundPixelShader();

		const FMaterialRenderProxy* MaterialProxyForRendering = MaterialProxy;
		const FMaterial& Material = MaterialProxy->GetMaterialWithFallback(View.GetFeatureLevel(), MaterialProxyForRendering);
		FMaterialShader::SetParameters(RHICmdList, ShaderRHI, MaterialProxyForRendering, Material, View);

		const FLargeWorldRenderPosition AbsoluteOrigin(View.ViewMatrices.GetInvViewMatrix().GetOrigin());
		const FVector3f TilePosition = AbsoluteOrigin.GetTile();
		const FMatrix WorldToDecalMatrix = DecalProxy.ComponentTrans.ToInverseMatrixWithScale();
		const FMatrix44f RelativeWorldToDecalMatrix = FLargeWorldRenderScalar::MakeFromRelativeWorldMatrix(AbsoluteOrigin.GetTileOffset(), WorldToDecalMatrix);
		const FMatrix DecalToWorldMatrix = DecalProxy.ComponentTrans.ToMatrixWithScale();
		const FMatrix44f RelativeDecalToWorldMatrix = FLargeWorldRenderScalar::MakeToRelativeWorldMatrix(AbsoluteOrigin.GetTileOffset(), DecalToWorldMatrix);
		const FVector3f OrientationVector = (FVector3f)DecalProxy.ComponentTrans.GetUnitAxis(EAxis::X);

		if (DecalTilePosition.IsBound())
		{
			SetShaderValue(RHICmdList, ShaderRHI, DecalTilePosition, TilePosition);
		}
		if(SvPositionToDecal.IsBound())
		{
			FVector2D InvViewSize = FVector2D(1.0f / View.ViewRect.Width(), 1.0f / View.ViewRect.Height());

			// setup a matrix to transform float4(SvPosition.xyz,1) directly to Decal (quality, performance as we don't need to convert or use interpolator)
			//	new_xy = (xy - ViewRectMin.xy) * ViewSizeAndInvSize.zw * float2(2,-2) + float2(-1, 1);
			//  transformed into one MAD:  new_xy = xy * ViewSizeAndInvSize.zw * float2(2,-2)      +       (-ViewRectMin.xy) * ViewSizeAndInvSize.zw * float2(2,-2) + float2(-1, 1);

			float Mx = 2.0f * InvViewSize.X;
			float My = -2.0f * InvViewSize.Y;
			float Ax = -1.0f - 2.0f * View.ViewRect.Min.X * InvViewSize.X;
			float Ay = 1.0f + 2.0f * View.ViewRect.Min.Y * InvViewSize.Y;

			// todo: we could use InvTranslatedViewProjectionMatrix and TranslatedWorldToComponent for better quality
			const FMatrix44f SvPositionToDecalValue = FMatrix44f(										// LWC_TODO: Precision loss
				FMatrix(
					FPlane(Mx,  0,   0,  0),
					FPlane( 0, My,   0,  0),
					FPlane( 0,  0,   1,  0),
					FPlane(Ax, Ay,   0,  1)
				) * View.ViewMatrices.GetInvViewProjectionMatrix() * WorldToDecalMatrix);

			SetShaderValue(RHICmdList, ShaderRHI, SvPositionToDecal, SvPositionToDecalValue);
		}
		if(DecalToWorld.IsBound())
		{
			SetShaderValue(RHICmdList, ShaderRHI, DecalToWorld, RelativeDecalToWorldMatrix);
		}
		if (DecalToWorldInvScale.IsBound())
		{
			SetShaderValue(RHICmdList, ShaderRHI, DecalToWorldInvScale, RelativeDecalToWorldMatrix.GetScaleVector().Reciprocal());
		}
		if (DecalOrientation.IsBound())
		{
			SetShaderValue(RHICmdList, ShaderRHI, DecalOrientation, OrientationVector);
		}
		
		float LifetimeAlpha = 1.0f;

		// Certain engine captures (e.g. environment reflection) don't have a tick. Default to fully opaque.
		if (View.Family->Time.GetWorldTimeSeconds())
		{
			LifetimeAlpha = FMath::Clamp(FMath::Min(View.Family->Time.GetWorldTimeSeconds() * -DecalProxy.InvFadeDuration + DecalProxy.FadeStartDelayNormalized, View.Family->Time.GetWorldTimeSeconds() * DecalProxy.InvFadeInDuration + DecalProxy.FadeInStartDelayNormalized), 0.0f, 1.0f);
		}
 
		SetShaderValue(RHICmdList, ShaderRHI, DecalParams, FVector2f(FadeAlphaValue, LifetimeAlpha));
	}

private:
	LAYOUT_FIELD(FShaderParameter, SvPositionToDecal);
	LAYOUT_FIELD(FShaderParameter, DecalTilePosition);
	LAYOUT_FIELD(FShaderParameter, DecalToWorld);
	LAYOUT_FIELD(FShaderParameter, DecalToWorldInvScale);
	LAYOUT_FIELD(FShaderParameter, DecalOrientation);
	LAYOUT_FIELD(FShaderParameter, DecalParams);
};

IMPLEMENT_MATERIAL_SHADER_TYPE(,FDeferredDecalPS,TEXT("/Engine/Private/DeferredDecal.usf"),TEXT("MainPS"),SF_Pixel);

class FDeferredDecalEmissivePS : public FDeferredDecalPS
{
	DECLARE_SHADER_TYPE(FDeferredDecalEmissivePS, Material);

public:
	static bool ShouldCompilePermutation(const FMaterialShaderPermutationParameters& Parameters)
	{
		return (Parameters.MaterialParameters.MaterialDomain == MD_DeferredDecal) &&
			DecalRendering::IsCompatibleWithRenderStage(DecalRendering::ComputeDecalBlendDesc(Parameters.Platform, Parameters.MaterialParameters), EDecalRenderStage::Emissive);
	}
	static void ModifyCompilationEnvironment(const FMaterialShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment)
	{
		FMaterialShader::ModifyCompilationEnvironment(Parameters, OutEnvironment);
		DecalRendering::ModifyCompilationEnvironment(DecalRendering::ComputeDecalBlendDesc(Parameters.Platform, Parameters.MaterialParameters), EDecalRenderStage::Emissive, OutEnvironment);
	}

	FDeferredDecalEmissivePS() {}
	FDeferredDecalEmissivePS(const ShaderMetaType::CompiledShaderInitializerType& Initializer)
		: FDeferredDecalPS(Initializer)
	{}
};

IMPLEMENT_MATERIAL_SHADER_TYPE(, FDeferredDecalEmissivePS, TEXT("/Engine/Private/DeferredDecal.usf"), TEXT("MainPS"), SF_Pixel);

class FDeferredDecalAmbientOcclusionPS : public FDeferredDecalPS
{
	DECLARE_SHADER_TYPE(FDeferredDecalAmbientOcclusionPS, Material);

public:
	static bool ShouldCompilePermutation(const FMaterialShaderPermutationParameters& Parameters)
	{
		return (Parameters.MaterialParameters.MaterialDomain == MD_DeferredDecal) &&
			DecalRendering::IsCompatibleWithRenderStage(DecalRendering::ComputeDecalBlendDesc(Parameters.Platform, Parameters.MaterialParameters), EDecalRenderStage::AmbientOcclusion);
	}
	static void ModifyCompilationEnvironment(const FMaterialShaderPermutationParameters& Parameters, FShaderCompilerEnvironment& OutEnvironment)
	{
		FMaterialShader::ModifyCompilationEnvironment(Parameters, OutEnvironment);
		DecalRendering::ModifyCompilationEnvironment(DecalRendering::ComputeDecalBlendDesc(Parameters.Platform, Parameters.MaterialParameters), EDecalRenderStage::AmbientOcclusion, OutEnvironment);
	}

	FDeferredDecalAmbientOcclusionPS() {}
	FDeferredDecalAmbientOcclusionPS(const ShaderMetaType::CompiledShaderInitializerType& Initializer)
		: FDeferredDecalPS(Initializer)
	{}
};

IMPLEMENT_MATERIAL_SHADER_TYPE(, FDeferredDecalAmbientOcclusionPS, TEXT("/Engine/Private/DeferredDecal.usf"), TEXT("MainPS"), SF_Pixel);

namespace DecalRendering
{
	float GetDecalFadeScreenSizeMultiplier()
	{
		return CVarDecalFadeScreenSizeMultiplier.GetValueOnRenderThread();
	}

	float CalculateDecalFadeAlpha(float DecalFadeScreenSize, const FMatrix& ComponentToWorldMatrix, const FViewInfo& View, float FadeMultiplier)
	{
		check(View.IsPerspectiveProjection());

		float Distance = (View.ViewMatrices.GetViewOrigin() - ComponentToWorldMatrix.GetOrigin()).Size();
		float Radius = ComponentToWorldMatrix.GetMaximumAxisScale();
		float CurrentScreenSize = ((Radius / Distance) * FadeMultiplier);

		// fading coefficient needs to increase with increasing field of view and decrease with increasing resolution
		// FadeCoeffScale is an empirically determined constant to bring us back roughly to fraction of screen size for FadeScreenSize
		const float FadeCoeffScale = 600.0f;
		float FOVFactor = ((2.0f / View.ViewMatrices.GetProjectionMatrix().M[0][0]) / View.ViewRect.Width()) * FadeCoeffScale;
		float FadeCoeff = DecalFadeScreenSize * FOVFactor;
		float FadeRange = FadeCoeff * 0.5f;

		float Alpha = (CurrentScreenSize - FadeCoeff) / FadeRange;
		return FMath::Clamp(Alpha, 0.0f, 1.0f);
	}

	void SortDecalList(FTransientDecalRenderDataList& Decals)
	{
		// Sort by sort order to allow control over composited result
		// Then sort decals by state to reduce render target switches
		// Also sort by component since Sort() is not stable
		struct FCompareFTransientDecalRenderData
		{
			FORCEINLINE bool operator()(const FTransientDecalRenderData& A, const FTransientDecalRenderData& B) const
			{
				if (B.Proxy.SortOrder != A.Proxy.SortOrder)
				{ 
					return A.Proxy.SortOrder < B.Proxy.SortOrder;
				}
				if (B.BlendDesc.bWriteNormal != A.BlendDesc.bWriteNormal)
				{
					// bWriteNormal here has priority because we want to render decals that output normals before those could read normals.
					// Also this is the only flag that can trigger a change of EDecalRenderTargetMode inside a single EDecalRenderStage, and we batch according to this.
					return B.BlendDesc.bWriteNormal < A.BlendDesc.bWriteNormal; // < so that those outputting normal are first.
				}
				if (B.BlendDesc.Packed != A.BlendDesc.Packed)
				{
					// Sorting by the FDecalBlendDesc contents will reduce blend state changes.
					return (int32)B.BlendDesc.Packed < (int32)A.BlendDesc.Packed;
				}
				if (B.MaterialProxy != A.MaterialProxy)
				{
					// Batch decals with the same material together
					return B.MaterialProxy < A.MaterialProxy;
				}
				return (PTRINT)B.Proxy.Component < (PTRINT)A.Proxy.Component;
			}
		};

		// Sort decals by blend mode to reduce render target switches
		Decals.Sort(FCompareFTransientDecalRenderData());
	}

	bool BuildVisibleDecalList(const FScene& Scene, const FViewInfo& View, EDecalRenderStage DecalRenderStage, FTransientDecalRenderDataList* OutVisibleDecals)
	{
		QUICK_SCOPE_CYCLE_COUNTER(BuildVisibleDecalList);

		if (OutVisibleDecals)
		{
			OutVisibleDecals->Empty(Scene.Decals.Num());
		}

		// Don't draw for shader complexity mode.
		// todo: Handle shader complexity mode for deferred decal.
		if (View.Family->EngineShowFlags.ShaderComplexity)
		{
			return false;
		}

		const float FadeMultiplier = GetDecalFadeScreenSizeMultiplier();
		const EShaderPlatform ShaderPlatform = View.GetShaderPlatform();

		const bool bIsPerspectiveProjection = View.IsPerspectiveProjection();

		// Build a list of decals that need to be rendered for this view in SortedDecals
		for (const FDeferredDecalProxy* DecalProxy : Scene.Decals)
		{
			if (!DecalProxy->DecalMaterial || !DecalProxy->DecalMaterial->IsValidLowLevelFast())
			{
				continue;
			}

			bool bIsShown = true;

			if (!DecalProxy->IsShown(&View))
			{
				bIsShown = false;
			}

			const FMatrix ComponentToWorldMatrix = DecalProxy->ComponentTrans.ToMatrixWithScale();

			// can be optimized as we test against a sphere around the box instead of the box itself
			const float ConservativeRadius = FMath::Sqrt(
					ComponentToWorldMatrix.GetScaledAxis(EAxis::X).SizeSquared() +
					ComponentToWorldMatrix.GetScaledAxis(EAxis::Y).SizeSquared() +
					ComponentToWorldMatrix.GetScaledAxis(EAxis::Z).SizeSquared());

			// can be optimized as the test is too conservative (sphere instead of OBB)
			if(ConservativeRadius < SMALL_NUMBER || !View.ViewFrustum.IntersectSphere(ComponentToWorldMatrix.GetOrigin(), ConservativeRadius))
			{
				bIsShown = false;
			}

			if (bIsShown)
			{
				FTransientDecalRenderData Data(Scene, *DecalProxy, ConservativeRadius);
			
				if (IsCompatibleWithRenderStage(Data.BlendDesc, DecalRenderStage))
				{
					if (bIsPerspectiveProjection && Data.Proxy.FadeScreenSize != 0.0f)
					{
						Data.FadeAlpha = CalculateDecalFadeAlpha(Data.Proxy.FadeScreenSize, ComponentToWorldMatrix, View, FadeMultiplier);
					}

					const bool bShouldRender = Data.FadeAlpha > 0.0f;

					if (bShouldRender)
					{
						if (!OutVisibleDecals)
						{
							return true;
						}
						OutVisibleDecals->Add(Data);
					}
				}
			}
		}

		if (!OutVisibleDecals)
		{
			return false;
		}

		if (OutVisibleDecals->Num() > 0)
		{
			SortDecalList(*OutVisibleDecals);

			return true;
		}

		return false;
	}

	FMatrix ComputeComponentToClipMatrix(const FViewInfo& View, const FMatrix& DecalComponentToWorld)
	{
		FMatrix ComponentToWorldMatrixTrans = DecalComponentToWorld.ConcatTranslation(View.ViewMatrices.GetPreViewTranslation());
		return ComponentToWorldMatrixTrans * View.ViewMatrices.GetTranslatedViewProjectionMatrix();
	}

	void SetShader(FRHICommandList& RHICmdList, FGraphicsPipelineStateInitializer& GraphicsPSOInit, uint32 StencilRef, const FViewInfo& View,
		const FTransientDecalRenderData& DecalData, EDecalRenderStage DecalRenderStage, const FMatrix& FrustumComponentToClip)
	{
		const FMaterialShaderMap* MaterialShaderMap = DecalData.MaterialResource->GetRenderingThreadShaderMap();
		const EDebugViewShaderMode DebugViewMode = View.Family->GetDebugViewShaderMode();

		// When in shader complexity, decals get rendered as emissive even though there might not be emissive decals.
		// FDeferredDecalEmissivePS might not be available depending on the decal blend mode.
		TShaderRef<FDeferredDecalPS> PixelShader;
		if (DecalRenderStage == EDecalRenderStage::Emissive || DebugViewMode != DVSM_None)
		{
			PixelShader = TShaderRef<FDeferredDecalPS>(MaterialShaderMap->GetShader<FDeferredDecalEmissivePS>());
		}
		else if (DecalRenderStage == EDecalRenderStage::AmbientOcclusion)
		{
			PixelShader = TShaderRef<FDeferredDecalPS>(MaterialShaderMap->GetShader<FDeferredDecalAmbientOcclusionPS>());
		}
		else
		{
			PixelShader = MaterialShaderMap->GetShader<FDeferredDecalPS>();
		}
		check(!PixelShader.IsNull());

		TShaderMapRef<FDeferredDecalVS> VertexShader(View.ShaderMap);

		{
			GraphicsPSOInit.BoundShaderState.VertexDeclarationRHI = GetVertexDeclarationFVector4();
			GraphicsPSOInit.BoundShaderState.VertexShaderRHI = VertexShader.GetVertexShader();
			GraphicsPSOInit.BoundShaderState.PixelShaderRHI = PixelShader.GetPixelShader();
			GraphicsPSOInit.PrimitiveType = PT_TriangleList;

			SetGraphicsPipelineState(RHICmdList, GraphicsPSOInit, StencilRef);
		}

		// Set vertex shader parameters.
		{
			FDeferredDecalVS::FParameters ShaderParameters;
			ShaderParameters.FrustumComponentToClip = FMatrix44f(FrustumComponentToClip);			// LWC_TODO: Precision loss?
			ShaderParameters.PrimitiveUniformBuffer = GIdentityPrimitiveUniformBuffer.GetUniformBufferRef();

			SetShaderParameters(RHICmdList, VertexShader, VertexShader.GetVertexShader(), ShaderParameters);
		}

		// Set pixel shader parameters.
		{
			PixelShader->SetParameters(RHICmdList, View, DecalData.MaterialProxy, DecalData.Proxy, DecalData.FadeAlpha);

			auto& PrimitivePS = PixelShader->GetUniformBufferParameter<FPrimitiveUniformShaderParameters>();

			// uncomment to track down usage of the Primitive uniform buffer
			//	check(!PrimitiveVS.IsBound());
			//	check(!PrimitivePS.IsBound());

			if (DebugViewMode == DVSM_None)
			{
				SetUniformBufferParameter(RHICmdList, PixelShader.GetPixelShader(), PrimitivePS, GIdentityPrimitiveUniformBuffer);
			}
		}

		// Set stream source after updating cached strides
		RHICmdList.SetStreamSource(0, GetUnitCubeVertexBuffer(), 0);
	}

	void SetVertexShaderOnly(FRHICommandList& RHICmdList, FGraphicsPipelineStateInitializer& GraphicsPSOInit, const FViewInfo& View, const FMatrix& FrustumComponentToClip)
	{
		TShaderMapRef<FDeferredDecalVS> VertexShader(View.ShaderMap);

		GraphicsPSOInit.BoundShaderState.VertexDeclarationRHI = GetVertexDeclarationFVector4();
		GraphicsPSOInit.BoundShaderState.VertexShaderRHI = VertexShader.GetVertexShader();
		GraphicsPSOInit.PrimitiveType = PT_TriangleList;

		SetGraphicsPipelineState(RHICmdList, GraphicsPSOInit, 0);

		// Set vertex shader parameters.
		{
			FDeferredDecalVS::FParameters ShaderParameters;
			ShaderParameters.FrustumComponentToClip = FMatrix44f(FrustumComponentToClip);	// LWC_TODO: Precision loss
			ShaderParameters.PrimitiveUniformBuffer = GIdentityPrimitiveUniformBuffer.GetUniformBufferRef();

			SetShaderParameters(RHICmdList, VertexShader, VertexShader.GetVertexShader(), ShaderParameters);
		}

	}
}