/*
* Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
*
* NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
* property and proprietary rights in and to this material, related
* documentation and any modifications thereto. Any use, reproduction,
* disclosure or distribution of this material and related documentation
* without an express license agreement from NVIDIA CORPORATION or
* its affiliates is strictly prohibited.
*/


#include "/Engine/Private/Common.ush"
#include "/Engine/Private/DeferredShadingCommon.ush"
#include "/Engine/Private/BRDF.ush"
#include "/Engine/Private/FastMath.ush"
#include "/Engine/Private/ScreenPass.ush"
#include "/Engine/Private/RayTracing/RayTracingCommon.ush"


#ifndef DIFFUSE_SPECULAR_ALBEDO
#define DIFFUSE_SPECULAR_ALBEDO 0
#endif

#ifndef SPECULAR_HITT
#define SPECULAR_HITT 0
#endif

SCREEN_PASS_TEXTURE_VIEWPORT(InputViewPort)
SCREEN_PASS_TEXTURE_VIEWPORT(OutputViewPort)


#include "/Engine/Private/PositionReconstructionCommon.ush"

Texture2D<float4> PassthroughDiffuse;
Texture2D<float4> PassthroughSpecular;
Texture2D<float4> PassthroughNormalRoughness;
Texture2D<float> PassthroughDepth;
Texture2D<float> ReflectionHitDistance;

float3 EnvBRDFApproxRTG(float3 SpecularColor, float alpha, float NoV)
{
	NoV = abs(NoV);
	// [Ray Tracing Gems, Chapter 32]
	float4 X;
	X.x = 1.f;
	X.y = NoV;
	X.z = NoV * NoV;
	X.w = NoV * X.z;

	float4 Y;
	Y.x = 1.f;
	Y.y = alpha;
	Y.z = alpha * alpha;
	Y.w = alpha * Y.z;

	float2x2 M1 = float2x2(0.99044f, -1.28514f, 1.29678f, -0.755907f);
	float3x3 M2 = float3x3(1.f, 2.92338f, 59.4188f, 20.3225f, -27.0302f, 222.592f, 121.563f, 626.13f, 316.627f);

	float2x2 M3 = float2x2(0.0365463f, 3.32707, 9.0632f, -9.04756);
	float3x3 M4 = float3x3(1.f, 3.59685f, -1.36772f, 9.04401f, -16.3174f, 9.22949f, 5.56589f, 19.7886f, -20.2123f);

	float bias = dot(mul(M1, X.xy), Y.xy) * rcp(dot(mul(M2, X.xyw), Y.xyw));
	float scale = dot(mul(M3, X.xy), Y.xy) * rcp(dot(mul(M4, X.xzw), Y.xyw));

	// This is a hack for specular reflectance of 0
	bias *= saturate(SpecularColor.g * 50);

	return mad(SpecularColor, max(0, scale), max(0, bias));
}

void GBufferResolvePixelShader(
	float2 InUV : TEXCOORD0,
	float4 SvPosition : SV_Position
#if DIFFUSE_SPECULAR_ALBEDO
	, out float4 OutDiffuseAlbedo : SV_Target0
	, out float4 OutSpecularAlbedo : SV_Target1
	, out float4 OutNormal : SV_Target2
	, out float OutRoughness : SV_Target3
	, out float OutDepth : SV_Target4
#endif
#if SPECULAR_HITT
	, out float OutHitT : SV_Target5
#endif
)
{
	uint2 PixelCoord = uint2(SvPosition.xy + InputViewPort_ViewportMin);


#if PASSTHROUGH_FEATURE_BUFFERS
	// In this case, everything has been pre-composited

#if DIFFUSE_SPECULAR_ALBEDO
	const float4 GuideDiffuse = PassthroughDiffuse[PixelCoord];
	const float4 GuideSpecular = PassthroughSpecular[PixelCoord];
	const float4 GuideNormalRoughness = PassthroughNormalRoughness[PixelCoord];
	const float GuideDepthData = PassthroughDepth[PixelCoord];

	OutDiffuseAlbedo = float4(GuideDiffuse.xyz, 1.0f);
	OutSpecularAlbedo = float4(GuideSpecular.xyz, 1.0f);
	OutNormal = float4(GuideNormalRoughness.xyz * 2.0 - 1.0, GuideNormalRoughness.w);
	OutRoughness = GuideNormalRoughness.w;

	OutDepth = GuideDepthData;
#endif

#if SPECULAR_HITT
	float HitT = ReflectionHitDistance.Load(int3(SvPosition.xy + InputViewPort_ViewportMin, 0));

	OutHitT = HitT;
#endif

#else
	const float DeviceZ = SceneTexturesStruct.SceneDepthTexture[PixelCoord].x;

	if (DeviceZ == 0.f)
	{
#if DIFFUSE_SPECULAR_ALBEDO
		const float2 UV = SvPosition.xy * View.BufferSizeAndInvSize.zw;
		const FRayDesc Ray = CreatePrimaryRay(UV);
		OutDiffuseAlbedo = 0.2f;
		OutSpecularAlbedo = 0.2f;
		OutNormal = float4(-Ray.Direction, 0.5f);
		OutRoughness = 0.5f;
		OutDepth = POSITIVE_INFINITY;
#endif

#if SPECULAR_HITT
		OutHitT = MaxHalfFloat;
#endif
		return;
	}

	FGBufferData GBuffer = GetGBufferDataUint(PixelCoord, true);

#if DIFFUSE_SPECULAR_ALBEDO

	float3 DiffuseColor = GBuffer.DiffuseColor;
	float3 SpecularColor = GBuffer.SpecularColor;

	if (UseSubsurfaceProfile(GBuffer.ShadingModelID))
	{
		// this is a bit odd, but the diffuse signal is multiplied by this without any conversion tied to metalness
		// See apply code in PostprocessSubsurface
		DiffuseColor = GBuffer.StoredBaseColor;

		// the correct value here is a bit debatable, as the skin computes with one value, then applies with another
		SpecularColor = ComputeF0(GBuffer.Specular, GBuffer.StoredBaseColor, GBuffer.Metallic);
		//SpecularColor = 1.0;
	}

	float3 WorldPos = ReconstructWorldPositionFromDepth(ViewportUVToBufferUV(InUV), GBuffer.Depth);
	float3 RayDirection = normalize(WorldPos - LWCHackToFloat(PrimaryView.WorldCameraOrigin));

	float NoV = saturate(dot(-RayDirection, GBuffer.WorldNormal));
	float x = NoV * NoV;
	NoV = x / (x + (1 - NoV) * (1 - NoV));
	float3 F = EnvBRDFApproxRTG(SpecularColor, Square(GBuffer.Roughness), NoV * 0.75f);

	OutDiffuseAlbedo = float4(DiffuseColor, 1.0f);
	OutSpecularAlbedo = float4(F, 1.0f);
	OutNormal = float4(GBuffer.WorldNormal, GBuffer.Roughness);
	OutRoughness = GBuffer.Roughness;

	OutDepth = GBuffer.Depth;
#endif

#if SPECULAR_HITT
	float HitT = ReflectionHitDistance.Load(int3(SvPosition.xy + InputViewPort_ViewportMin, 0));

	OutHitT = HitT;
#endif
#endif
}