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

#pragma once

#if SUBSTRATE_ENABLED && SUBSTRATE_GBUFFER_FORMAT==0 
#define SUBSTRATE_LOAD_FROM_MATERIALCONTAINER 0
#endif

#include "../DeferredShadingCommon.ush"
#include "../Lumen/LumenPosition.ush"

#include "../HairStrands/HairStrandsVisibilityCommon.ush"
#include "../HairStrands/HairStrandsVisibilityUtils.ush"
#include "../HairStrands/HairStrandsCommon.ush"
#include "../HairStrands/HairStrandsDeepTransmittanceCommon.ush"
#include "../HairStrands/HairStrandsDeepTransmittanceDualScattering.ush"

#include "../Substrate/Substrate.ush"
#include "../Substrate/SubstrateEvaluation.ush"
#include "../Substrate/SubstrateDeferredLighting.ush"
#include "../Substrate/SubstrateMaterialSampling.ush"
#include "../Substrate/SubstrateTile.ush"
#include "../Substrate/SubstrateRead.ush"

// Types of FMegaLightsMaterial input
#define INPUT_TYPE_GBUFFER 0
#define INPUT_TYPE_HAIRSTRANDS 1

#ifndef INPUT_TYPE
	#define INPUT_TYPE INPUT_TYPE_GBUFFER
#endif

struct FMegaLightsMaterial
{
	float  Depth;
	float3 WorldNormal;
	float3 WorldNormalForPositionBias;
	float  Roughness;
	bool   bIsValid;
	bool   bIsSimple;
	bool   bIsSingle;
	bool   bIsComplexSpecial;
	bool   bIsHair;
	bool   bHasBackfaceDiffuse;
	bool   bNeedsSeparateSubsurfaceLightAccumulation;
	bool   bNeedsComplexTransmittance;
	bool   bAllowSpatialFilter;
	uint   LightingChannelMask; // Not loaded by default
	float  Coverage;
	uint   GlobalIndex1D;
	uint2  GlobalIndex2D;

	float3 DiffuseColor;
	float3 SpecularColor;

	bool IsValid()
	{
		return bIsValid;
	}

	bool IsSimple()
	{
		return bIsSimple;
	}

	bool IsSingle()
	{
		return bIsSingle;
	}

	bool IsComplexSpecial()
	{
		return bIsComplexSpecial;
	}

	void SetDepth(float In)
	{
		Depth = In;
		#if !SUBSTRATE_ENABLED || SUBSTRATE_GBUFFER_FORMAT==0
		GBuffer.Depth = In;
		#endif
	}

	#if INPUT_TYPE == INPUT_TYPE_HAIRSTRANDS
	FGBufferData GBuffer;
	uint MacroGroupId;
	#elif SUBSTRATE_ENABLED && SUBSTRATE_GBUFFER_FORMAT==1
	uint ClosureIndex;
	float PDF;
	FSubstrateBSDF BSDF;
	float3x3 BSDFTangentBasis;
	float BSDFAO;
	#elif SUBSTRATE_ENABLED && SUBSTRATE_GBUFFER_FORMAT==0
	FGBufferData GBuffer;
	FSubstrateBSDF BSDF;
	float3x3 BSDFTangentBasis;
	float BSDFAO;
	#else
	FGBufferData GBuffer;
	#endif
};

FMegaLightsMaterial LoadMaterial(float2 ScreenUV, uint2 ScreenCoord, bool bForceSimpleShading = false, uint SampleIt=0)
#if INPUT_TYPE == INPUT_TYPE_HAIRSTRANDS
{
	const uint3 PixelCoord = uint3(ScreenCoord, 0);
	const float HairDeviceZ = HairStrands.HairOnlyDepthTexture.Load(PixelCoord).x;
	const FNodeDesc NodeDesc = DecodeNodeDesc(HairStrands.HairSampleOffset.Load(PixelCoord));
	const float PixelCoverage = min(HairStrands.HairCoverageTexture.Load(PixelCoord), 1);

	FMegaLightsMaterial Out = (FMegaLightsMaterial)0;
	Out.Depth = ConvertFromDeviceZ(HairDeviceZ);
	Out.bIsValid = HairDeviceZ > 0 && NodeDesc.Count > 0 && PixelCoverage > 0 && SampleIt < NodeDesc.Count;

	const uint TotalNodeCount = HairStrands.HairSampleCount[0];
	const uint2 Resolution = GetHairSampleResolution(TotalNodeCount);

	if (Out.bIsValid)
	{
		{
			const FHairSample Sample = UnpackHairSample(HairStrands.HairSampleData[NodeDesc.Offset + SampleIt]);
			FGBufferData GBuffer = HairSampleToGBufferData(Sample);
			const bool bIsLit = GBuffer.ShadingModelID != SHADINGMODELID_UNLIT;
			GBuffer.Roughness = max(GBuffer.Roughness, View.MinRoughness);

			Out.bIsValid   		= bIsLit;
			Out.bIsSimple  		= false; //GBuffer.ShadingModelID == SHADINGMODELID_DEFAULT_LIT;
			Out.bIsSingle		= false;
			Out.bIsComplexSpecial= false;
			Out.bIsHair 		= true;
			Out.bHasBackfaceDiffuse = false;
			Out.bNeedsSeparateSubsurfaceLightAccumulation = false;
			Out.bNeedsComplexTransmittance = true;
			Out.WorldNormal  	= GBuffer.WorldNormal;
			Out.WorldNormalForPositionBias = GBuffer.WorldNormal;
			Out.Depth        	= GBuffer.Depth;
			Out.Roughness    	= GBuffer.Roughness;
			Out.DiffuseColor 	= GBuffer.DiffuseColor;
			Out.SpecularColor	= GBuffer.SpecularColor;
			Out.GBuffer      	= GBuffer;
			Out.bAllowSpatialFilter = true; // false ?
			{
				Out.WorldNormalForPositionBias = ComputeHairNormalForPositionBias(ScreenUV, Out.Depth, Out.WorldNormal);
			}

			Out.MacroGroupId  = Sample.MacroGroupId;
			Out.LightingChannelMask = Sample.LightChannelMask;
			Out.Coverage = saturate(From8bitCoverage(Sample.Coverage8bit));
			Out.GlobalIndex1D = NodeDesc.Offset + SampleIt;
			Out.GlobalIndex2D = GetHairSampleCoord(Out.GlobalIndex1D, Resolution);
		}
	}
	return Out;
}
#elif INPUT_TYPE == INPUT_TYPE_GBUFFER && SUBSTRATE_GBUFFER_FORMAT==1
{
	#if SUBSTRATE_STOCHASTIC_LIGHTING
	if (Substrate.bStochasticLighting)
	{
		const FSubstrateSampledMaterial SampledMaterial = UnpackSampledMaterial(Substrate.SampledMaterialTexture[ScreenCoord]);

		FMegaLightsMaterial Out = (FMegaLightsMaterial)0;
		Out.WorldNormal = SampledMaterial.WorldNormal;
		Out.WorldNormalForPositionBias = SampledMaterial.WorldNormal;
		Out.Depth = ConvertFromDeviceZ(SceneTexturesStruct.SceneDepthTexture.Load(int3(ScreenCoord, 0)).r);
		Out.Roughness = SampledMaterial.Roughness;
		Out.bAllowSpatialFilter = SampledMaterial.bAllowSpatialFilter;

		Out.bIsValid = SampledMaterial.bIsValid;
		Out.bIsSimple = SampledMaterial.bIsSimple;
		Out.bIsSingle = SampledMaterial.bIsSingle;
		Out.bIsComplexSpecial = false;
		Out.bIsHair = SampledMaterial.bIsHair;

		// SUBSTRATE_TODO: implement Out.bHasBackfaceDiffuse
		Out.bHasBackfaceDiffuse	= false;
		Out.bNeedsComplexTransmittance = false;

		Out.bNeedsSeparateSubsurfaceLightAccumulation = SampledMaterial.bNeedsSeparateSubsurfaceLightAccumulation;
		Out.DiffuseColor = SampledMaterial.DiffuseAlbedo;
		Out.SpecularColor = SampledMaterial.SpecularAlbedo;
		Out.ClosureIndex = SampledMaterial.ClosureIndex;
		Out.PDF = SampledMaterial.PDF;
		Out.Coverage = 1;
		return Out;
	}
	else
	#endif // SUBSTRATE_STOCHASTIC_LIGHTING
	{
		const FSubstrateTopLayerData TopLayerData = SubstrateUnpackTopLayerData(Substrate.TopLayerTexture.Load(uint3(ScreenCoord, 0)));
		FMegaLightsMaterial Out = (FMegaLightsMaterial)0;
		Out.WorldNormal = TopLayerData.WorldNormal;
		Out.WorldNormalForPositionBias = TopLayerData.WorldNormal;
		Out.Depth = ConvertFromDeviceZ(SceneTexturesStruct.SceneDepthTexture.Load(int3(ScreenCoord, 0)).r);
		Out.Roughness = TopLayerData.Roughness;
		Out.bAllowSpatialFilter = true;
		Out.ClosureIndex = 0;
		Out.PDF = 1.f;
		Out.Coverage = 1;
#if SUBSTRATE_MATERIALCONTAINER_IS_VIEWRESOURCE

		const float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, Out.Depth);
		const float3 V = -GetCameraVectorFromTranslatedWorldPosition(TranslatedWorldPosition);

		const FSubstrateIntegrationSettings Settings = InitSubstrateIntegrationSettings(false /*bForceFullyRough*/, Substrate.bRoughDiffuse, Substrate.PeelLayersAboveDepth, Substrate.bRoughnessTracking);
		FSubstrateAddressing SubstrateAddressing = GetSubstratePixelDataByteOffset(ScreenCoord, uint2(View.BufferSizeAndInvSize.xy), Substrate.MaxBytesPerPixel);
		FSubstratePixelHeader SubstratePixelHeader = UnpackSubstrateHeaderIn(Substrate.MaterialTextureArray, SubstrateAddressing, Substrate.TopLayerTexture);

		Out.bIsValid = SubstratePixelHeader.IsSubstrateMaterial();
		Out.bIsSimple = SubstratePixelHeader.IsSimpleMaterial();
		Out.bIsSingle = SubstratePixelHeader.IsSingleMaterial();
		Out.bIsComplexSpecial = SubstratePixelHeader.IsComplexSpecialMaterial();
		Out.bIsHair = SubstratePixelHeader.IsHair();
		Out.bHasBackfaceDiffuse = SubstratePixelHeader.SubstrateGetBSDFType() == SUBSTRATE_BSDF_TYPE_SLAB && SubstratePixelHeader.HasSubsurface();
		Out.bNeedsComplexTransmittance = false;

		#if SUBSTRATE_LOAD_FROM_MATERIALCONTAINER
		Substrate_for(uint ClosureIndex = 0, ClosureIndex < SubstratePixelHeader.ClosureCount, ++ClosureIndex)
		#else
		if (SubstratePixelHeader.ClosureCount > 0)
		#endif
		{
			FSubstrateBSDF BSDF = UnpackSubstrateBSDF(Substrate.MaterialTextureArray, SubstrateAddressing, SubstratePixelHeader);

			// Create the BSDF context
			FSubstrateBSDFContext SubstrateBSDFContext = SubstrateCreateBSDFContext(SubstratePixelHeader, BSDF, SubstrateAddressing, V);
			const float3 BSDFThroughput = LuminanceWeight(SubstrateBSDFContext, BSDF);

			// Evaluate environment lighting
			FSubstrateEnvLightResult SubstrateEnvLight = SubstrateEvaluateForEnvLight(SubstrateBSDFContext, true /*bEnableSpecular*/, Settings);

			Out.DiffuseColor += BSDFThroughput * SubstrateEnvLight.DiffuseColor;  //SubstrateEnvLight.DiffuseWeight;
			Out.SpecularColor += BSDFThroughput * SubstrateEnvLight.SpecularColor; //SubstrateEnvLight.SpecularWeight;

			#if SUBSTRATE_FASTPATH==0
			if (any(SubstrateEnvLight.SpecularHazeWeight > 0.0f))
			{
				Out.SpecularColor += BSDFThroughput * SubstrateEnvLight.SpecularHazeWeight;
			}
			#endif 

			if (SubstrateEnvLight.bPostProcessSubsurface)
			{
				Out.bNeedsSeparateSubsurfaceLightAccumulation = true;
			}

			#if SUBSTRATE_FASTPATH==0
			if (SubstratePixelHeader.IsComplexSpecialMaterial() && BSDF_GETHASGLINT(BSDF))
			{
				Out.bAllowSpatialFilter = false;
			}
			#endif

			#if SUBSTRATE_LOAD_FROM_MATERIALCONTAINER == 0
			Out.BSDF = BSDF;
			Out.BSDFTangentBasis = SubstrateBSDFContext.TangentBasis;
			Out.BSDFAO = SubstrateGetAO(SubstratePixelHeader);
			#endif
		}
		#endif //SUBSTRATE_MATERIALCONTAINER_IS_VIEWRESOURCE

		return Out;
	}
}
#elif INPUT_TYPE == INPUT_TYPE_GBUFFER
{
	FGBufferData GBuffer = GetGBufferData(ScreenUV);
	GBuffer.Roughness = max(GBuffer.Roughness, View.MinRoughness);
	const bool bIsLit = GBuffer.ShadingModelID != SHADINGMODELID_UNLIT;

	if (bForceSimpleShading)
	{
		checkSlow(GBuffer.ShadingModelID == SHADINGMODELID_DEFAULT_LIT || GBuffer.ShadingModelID == SHADINGMODELID_UNLIT);
		GBuffer.ShadingModelID = SHADINGMODELID_DEFAULT_LIT;
	}

	FMegaLightsMaterial Out = (FMegaLightsMaterial)0;
	Out.bIsValid = bIsLit;
	Out.bIsSimple = GBuffer.ShadingModelID == SHADINGMODELID_DEFAULT_LIT;
	Out.bIsSingle = false;
	Out.bIsComplexSpecial = false;
	Out.bIsHair = GBuffer.ShadingModelID == SHADINGMODELID_HAIR;
	Out.bHasBackfaceDiffuse = GBuffer.ShadingModelID == SHADINGMODELID_TWOSIDED_FOLIAGE || GBuffer.ShadingModelID == SHADINGMODELID_SUBSURFACE;
	Out.bNeedsSeparateSubsurfaceLightAccumulation = UseSubsurfaceProfile(GBuffer.ShadingModelID);
	Out.bNeedsComplexTransmittance = GBuffer.ShadingModelID == SHADINGMODELID_HAIR && ShouldUseHairComplexTransmittance(GBuffer);
	Out.WorldNormal = GBuffer.WorldNormal;
	Out.WorldNormalForPositionBias = GBuffer.WorldNormal;
	Out.Depth = GBuffer.Depth;
	Out.Roughness = GBuffer.Roughness;
	Out.DiffuseColor = GBuffer.DiffuseColor;
	Out.SpecularColor = GBuffer.SpecularColor;
	Out.GBuffer = GBuffer;
	Out.bAllowSpatialFilter = true;
	Out.Coverage = 1;

	if (GBuffer.ShadingModelID == SHADINGMODELID_HAIR)
	{
		Out.WorldNormalForPositionBias = ComputeHairNormalForPositionBias(ScreenUV, Out.Depth, Out.WorldNormal);
	}

	#if SUBSTRATE_ENABLED && SUBSTRATE_GBUFFER_FORMAT==0

	const bool bIsComplex = abs(GBuffer.Anisotropy) > 0 || GBuffer.ShadingModelID == SHADINGMODELID_EYE || GBuffer.ShadingModelID == SHADINGMODELID_HAIR;
	Out.bIsSimple = Out.bIsSimple && abs(GBuffer.Anisotropy) == 0;
	Out.bIsSingle = !Out.bIsSimple && !bIsComplex;

	FScreenSpaceData ScreenSpaceData;
	ScreenSpaceData.GBuffer = GBuffer;
	ScreenSpaceData.AmbientOcclusion = 1.f;
	FSubstrateGBufferBSDF GBufferBSDF = SubstrateReadGBufferBSDF(ScreenSpaceData);
	Out.BSDF = GBufferBSDF.BSDF;
	Out.BSDFTangentBasis = GBufferBSDF.BSDFTangentBasis;
	Out.BSDFAO = GBufferBSDF.BSDFAO;
	#endif

	return Out;
}
#endif

struct FDenoisingModulateFactors
{
	float3 Diffuse;
	float3 Specular;
};

/**
 * Factors used to remove non-stochastic material detail before denoising in order to preserve texture detail
 * Divide before denoising and multiply after
 */ 
FDenoisingModulateFactors GetDenoisingModulateFactors(FMegaLightsMaterial Material, float3 TranslatedWorldPosition)
{
	const float3 N = Material.WorldNormal;
	const float3 V = normalize(View.TranslatedWorldCameraOrigin - TranslatedWorldPosition);
	const float NoV = saturate(dot(N, V));
	float3 SpecularEnv = EnvBRDF(Material.SpecularColor, Material.Roughness, NoV);
	
	// Hair should technically use the follow function, but it is expansive and does not improve visuals/stability. 
	#if 0
	if (Material.bIsHair)
	{
		float3 L = 0;
		SpecularEnv = EvaluateEnvHair(Material.GBuffer, V, N, L /*out*/);
	}
	#endif

	FDenoisingModulateFactors Factors;
	Factors.Diffuse = clamp(Material.DiffuseColor, 0.04f, 1.0f);
	Factors.Specular = clamp(SpecularEnv, 0.1f, 1.0f);
	return Factors;
}