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

#include "../Common.ush"

// ----------------------------------------------------------------------------------

#if defined (__INTELLISENSE__)
// Uncomment the appropriate define for enabling syntax highlighting with HLSL Tools for Visual Studio : 
//#define MERGE_EDIT_LAYERS 1
//#define FINALIZE_WEIGHTMAP 1
//#define GENERATE_MIPS 1
#endif // defined (__INTELLISENSE__)


// ----------------------------------------------------------------------------------

#if MERGE_EDIT_LAYERS
// MergeEditLayers inputs/outputs :
/** EEditLayerWeightmapBlendMode enum */
#define EEDITLAYERWEIGHTMAPBLENDMODE_ADDITIVE 0
#define EEDITLAYERWEIGHTMAPBLENDMODE_SUBSTRACTIVE 1

// Contains the information related to a given edit layers and how it should merge into the final weightmap
struct FEditLayerWeightmapMergeInfo
{
	uint SourceWeightmapTextureIndex; // The index in InPackedWeightmaps of the texture to read from for this layer
	uint SourceWeightmapTextureChannel; // The channel of the texture to read from for this layer
	uint BlendMode; // See EEditLayerWeightmapBlendMode
	float Alpha; // Alpha value to be used in the blend
};

uint InNumEditLayers; // Number of layers to merge into the destination texture
uint InStartIndexInEditLayersMergeInfos; // Start index for edit layers merge infos in the big buffer InEditLayersMergeInfos
Texture2DArray<float4> InPackedWeightmaps; // Source -packed- RGBA weightmaps
StructuredBuffer<FEditLayerWeightmapMergeInfo> InEditLayersMergeInfos; // Describes how to merge each individual edit layer (effective size == sum of InNumEditLayers of all passes)
#endif // MERGE_EDIT_LAYERS


#if FINALIZE_WEIGHTMAP
/** EWeightmapPaintLayerFlags enum */
#define EWEIGHTMAPPAINTLAYERFLAGS_ISVISIBILITYLAYER (1 << 0)
#define EWEIGHTMAPPAINTLAYERFLAGS_ISWEIGHTBLENDED (1 << 1)

// Contains the information related to a given edit layers and how it should merge into the final weightmap
struct FEditLayerWeightmapPaintLayerInfo
{
	uint Flags; // See EWeightmapPaintLayerFlags
};

// FinalizeWeightmap inputs/outputs :
uint InValidTextureChannelsMask; // 4 bits-wide bitmask that indicates which of the 4 channels need to be processed 

// Contains (per channel : .r for the first channel, .g for the 2nd, etc.) the index of the paint layer being processed for this channel. Note : this is an index for the channel's InPerChannelPaintLayerWeightmaps 
//  (which is a subset of all paint layers), not InPaintLayerInfos (i.e. all paint layers)
uint4 InPerChannelPaintLayerIndexInWeightmaps; 
// Contains (per channel : .r for the first channel, .g for the 2nd, etc.) the start index for addressing the indices of actual paint layers contained in the channel's InPerChannelPaintLayerWeightmaps
uint4 InPerChannelStartPaintLayerIndex; 
// Contains (per channel : .r for the first channel, .g for the 2nd, etc.) the number of actual paint layers contained in the channel's InPerChannelPaintLayerWeightmaps
uint4 InPerChannelNumPaintLayers; 

// 4 input texture arrays (R8) (one per channel) we're about to pack on the output weightmap (RGBA), each one contains a list of all paint layers (after having been edit layer-merged) for a given component, 
//  since we need to weight-blend some of them. For each channel, NumElements = InPerChannelNumPaintLayers[ChannelIndex]
Texture2DArray<float> InPerChannelPaintLayerWeightmaps_0; 
Texture2DArray<float> InPerChannelPaintLayerWeightmaps_1; 
Texture2DArray<float> InPerChannelPaintLayerWeightmaps_2; 
Texture2DArray<float> InPerChannelPaintLayerWeightmaps_3; 

// Input uint buffers that contains all the indices of the paint layer in InPaintLayerInfos. This buffer is shared between all texture resolve passes and is therefore adressed using 
//  InPerChannelStartPaintLayerIndex and InPerChannelNumPaintLayers (one entry per channel each)
StructuredBuffer<uint> InPaintLayerInfoIndices;

// Array that contains per-paint layer information (e.g. paint layer's blend mode)
StructuredBuffer<FEditLayerWeightmapPaintLayerInfo> InPaintLayerInfos; 
#endif // FINALIZE_WEIGHTMAP


#if GENERATE_MIPS
// GenerateMips inputs/outputs :
uint2 InCurrentMipSize; // Size of the texture mip currently being generated
uint InNumSubsections; // Number of sub-sections for this landscape (1 or 2)

Texture2D<float4> InSourceWeightmap;
#endif // GENERATE_MIPS


// ----------------------------------------------------------------------------------
// Util functions : 

#if MERGE_EDIT_LAYERS
// For a given edit layer, samples the appropriate texture in the input texture array and returns the appropriate channel, attenuated by the layer's alpha value : 
float SampleEditLayerWeightmap(uint2 InSampleCoordinates, uint InIndex)
{
	FEditLayerWeightmapMergeInfo LayerMergeInfo = InEditLayersMergeInfos[InStartIndexInEditLayersMergeInfos + InIndex];
	int4 LoadCoordinates = int4(InSampleCoordinates, LayerMergeInfo.SourceWeightmapTextureIndex, 0); // xy = relative coordinates, z = index in texture array, w = mip level
	float4 Packed = InPackedWeightmaps.Load(LoadCoordinates);
	float Weight = Packed[LayerMergeInfo.SourceWeightmapTextureChannel];
	return Weight * LayerMergeInfo.Alpha;
}
#endif // #if MERGE_EDIT_LAYERS

#if FINALIZE_WEIGHTMAP
float LoadPaintLayerWeightmapForChannel(uint InChannelIndex, uint InPaintLayerIndex, uint2 InSampleCoordinates)
{
	int4 LoadCoordinates = int4(InSampleCoordinates, InPaintLayerIndex, 0); // xy = coordinates, z = index in texture array, w = mip level
	// We have to use a switch case because we cannot sample a static array of textures without a literal expression:
	switch (InChannelIndex)
	{
		case 0: return InPerChannelPaintLayerWeightmaps_0.Load(LoadCoordinates);
		case 1: return InPerChannelPaintLayerWeightmaps_1.Load(LoadCoordinates);
		case 2: return InPerChannelPaintLayerWeightmaps_2.Load(LoadCoordinates);
		default: return InPerChannelPaintLayerWeightmaps_3.Load(LoadCoordinates);
	}
}

uint LoadPaintLayerInfoIndexForChannel(uint InChannelIndex, uint InPaintLayerIndex)
{
	uint PaintLayerInfoIndexInBigBuffer = InPerChannelStartPaintLayerIndex[InChannelIndex] + InPaintLayerIndex;
	return InPaintLayerInfoIndices[PaintLayerInfoIndexInBigBuffer];
}

bool IsWeightBlendedPaintLayer(FEditLayerWeightmapPaintLayerInfo InPaintLayerInfo)
{
	// For the visibility layer, deactivate weight blending altogether : 
	return (((InPaintLayerInfo.Flags & EWEIGHTMAPPAINTLAYERFLAGS_ISVISIBILITYLAYER) == 0)
		&& (InPaintLayerInfo.Flags & EWEIGHTMAPPAINTLAYERFLAGS_ISWEIGHTBLENDED));
}

float FinalizeWeightmapChannel(uint2 InSampleCoordinates, uint InChannelIndex)
{
	float Result = 0.0f;
	
	if (InValidTextureChannelsMask & (1u << InChannelIndex))
	{
		uint PaintLayerIndexInWeightmaps = InPerChannelPaintLayerIndexInWeightmaps[InChannelIndex];
		uint NumPaintLayers = InPerChannelNumPaintLayers[InChannelIndex];
		
		uint ActivePaintLayerInfoIndex = LoadPaintLayerInfoIndexForChannel(InChannelIndex, PaintLayerIndexInWeightmaps);
		FEditLayerWeightmapPaintLayerInfo ActivePaintLayerInfo = InPaintLayerInfos[ActivePaintLayerInfoIndex];

		if (!IsWeightBlendedPaintLayer(ActivePaintLayerInfo))
		{
			return LoadPaintLayerWeightmapForChannel(InChannelIndex, PaintLayerIndexInWeightmaps, InSampleCoordinates);
		}
		
		float ActivePaintLayerWeight = 0.0f;
		float BlendedWeightsSum = 0.0f;

		LOOP
		for (int i = 0; i < NumPaintLayers; ++i)
		{
			bool bIsOutputPaintLayer = (i == PaintLayerIndexInWeightmaps);
			uint PaintLayerInfoIndex = LoadPaintLayerInfoIndexForChannel(InChannelIndex, i);
			FEditLayerWeightmapPaintLayerInfo PaintLayerInfo = InPaintLayerInfos[PaintLayerInfoIndex];

			// Only weight blended (and non-visibility) paint layers participate to weight blending : 
			if (IsWeightBlendedPaintLayer(PaintLayerInfo)) 
			{
				float Weight = LoadPaintLayerWeightmapForChannel(InChannelIndex, i, InSampleCoordinates);
				if (bIsOutputPaintLayer)
				{
					ActivePaintLayerWeight = Weight;
				}
				BlendedWeightsSum += Weight;
			}
		}

		Result = (BlendedWeightsSum > 0.0f) ? saturate(ActivePaintLayerWeight / BlendedWeightsSum) : ActivePaintLayerWeight;
	}

	return Result;
}

#endif // FINALIZE_WEIGHTMAP


// ----------------------------------------------------------------------------------
// Pixel shaders : 

#if MERGE_EDIT_LAYERS
void MergeEditLayers(in float4 SVPos : SV_POSITION, out float OutColor : SV_Target0)
{
	uint2 TextureCoordinates = floor(SVPos.xy);
		
	// Take the base weight straight from the first layer : 
	float Weight = SampleEditLayerWeightmap(TextureCoordinates, 0);
		
	// Then merge the subsequent layers according to their merge info : 
	LOOP
	for (uint i = 1; i < InNumEditLayers; ++i)
	{
		FEditLayerWeightmapMergeInfo LayerMergeInfo = InEditLayersMergeInfos[InStartIndexInEditLayersMergeInfos + i];
		float LayerWeight = SampleEditLayerWeightmap(TextureCoordinates, i);
			
		if (LayerMergeInfo.BlendMode == EEDITLAYERWEIGHTMAPBLENDMODE_ADDITIVE)
		{
			Weight += LayerWeight;
		}
		else if (LayerMergeInfo.BlendMode == EEDITLAYERWEIGHTMAPBLENDMODE_SUBSTRACTIVE)
		{
			Weight -= LayerWeight;
		}
	}
	
	OutColor = Weight;
}
#endif // MERGE_EDIT_LAYERS


// ----------------------------------------------------------------------------------

#if FINALIZE_WEIGHTMAP
void FinalizeWeightmap(in float4 SVPos : SV_POSITION, out float4 OutColor : SV_Target0)
{
	uint2 TextureCoordinates = floor(SVPos.xy);
	OutColor = 0.0f;
	
	UNROLL
	for (uint i = 0; i < 4; ++i)
	{
		OutColor[i] = FinalizeWeightmapChannel(TextureCoordinates, i);
	}
}
#endif // FINALIZE_WEIGHTMAP


// ----------------------------------------------------------------------------------
#if GENERATE_MIPS
void GenerateMips(in float4 SVPos : SV_POSITION, out float4 OutColor : SV_Target0)
{
	uint2 TextureCoordinates = floor(SVPos.xy);
	
	float4 SourceSamples[4] =
	{
		InSourceWeightmap.Load(int3(2 * TextureCoordinates + int2(+0, +0), 0)),
		InSourceWeightmap.Load(int3(2 * TextureCoordinates + int2(+1, +0), 0)),
		InSourceWeightmap.Load(int3(2 * TextureCoordinates + int2(+0, +1), 0)),
		InSourceWeightmap.Load(int3(2 * TextureCoordinates + int2(+1, +1), 0)),
	};
		
	// Because the borders of each landscape components are shared between neighbors, we must ensure that the parent mip's inner row/column of pixels don't contribute : 
	// 9 possible cases (only the samples marked with a * must be kept): 
	//      bIsMinBorder.x = true
	//     |                        bIsMaxBorder.x = true
	//     |                       |
	//     v                       v
	// +-------+   +-------+   +-------+
	// | * :   |   | * : * |   |   : * |
	// | - + - |...| - + - |...| - + - | <-- bIsMinBorder.y = true
	// |   :   |   |   :   |   |   :   |
	// +-------+   +-------+   +-------+
	//    ...         ...	      ...
	// +-------+   +-------+   +-------+
	// | * :   |   | * : * |   |   : * |
	// | - + - |...| - + - |...| - + - |
	// | * :   |   | * : * |   |   : * |
	// +-------+   +-------+   +-------+
	//    ...         ...	      ...
	// +-------+   +-------+   +-------+
	// |   :   |   |   :   |   |   :   |
	// | - + - |...| - + - |...| - + - |
	// | * :   |   | * : * |   |   : * | <-- bIsMaxBorder.y = true
	// +-------+   +-------+   +-------+

	bool bIsLastMip = all(InCurrentMipSize == 1);

	bool2 bIsMinBorder = (TextureCoordinates == 0);
	bool2 bIsMaxBorder = (TextureCoordinates == (InCurrentMipSize - 1));
	// If we use subsections, we'll have an additional column/row of min border and another of max border next to them
	if (InNumSubsections == 2)
	{
		uint2 SubsectionSize = InCurrentMipSize / 2;
		bIsMinBorder = or(bIsMinBorder, TextureCoordinates == SubsectionSize);
		bIsMaxBorder = or(bIsMaxBorder, TextureCoordinates == SubsectionSize - 1);
	}
		
	float SampleWeights[4] =
	{ 
		// On the last mip, it's ok to keep all 4 samples : all neighbors components share them : 
		((bIsMaxBorder.x || bIsMaxBorder.y) && !bIsLastMip) ? 0.0f : 1.0f,
		((bIsMinBorder.x || bIsMaxBorder.y) && !bIsLastMip) ? 0.0f : 1.0f,
		((bIsMaxBorder.x || bIsMinBorder.y) && !bIsLastMip) ? 0.0f : 1.0f,
		((bIsMinBorder.x || bIsMinBorder.y) && !bIsLastMip) ? 0.0f : 1.0f,
	};
	
	float TotalSampleWeight = 0.0f;
	OutColor = 0.0f;
	
	UNROLL
	for (int i = 0; i < 4; ++i)
	{
		OutColor += SourceSamples[i] * SampleWeights[i];
		TotalSampleWeight += SampleWeights[i];
	}
	OutColor /= TotalSampleWeight;
}
#endif // GENERATE_MIPS