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

// When loading SSS checkerboard pixel, do not adjust DiffuseColor/SpecularColor to preserve specular and diffuse lighting values for each pixel
#define ALLOW_SSS_MATERIAL_OVERRIDE 0

#include "../Common.ush"
#include "MegaLights.ush"
#include "MegaLightsMaterial.ush"
#include "MegaLightsTileClassification.ush"

#ifdef MegaLightsTileClassificationBuildListsCS

Texture2D<uint> MegaLightsTileBitmask;

RWStructuredBuffer<uint> RWTileAllocator;
RWStructuredBuffer<uint> RWTileData;

uint2 ViewSizeInTiles;
uint2 ViewMinInTiles;
uint2 DownsampledViewMinInTiles;
uint2 DownsampledViewSizeInTiles;
uint OutputTileDataStride;

uint BuildTileMode(uint TileBitmask, uint InShading, uint InShadingRect, uint InShadingRectTextured)
{
	uint Out;
	Out = InShading;
	if (TileBitmask & MEGALIGHTS_TILE_BITMASK_RECT_LIGHT)
	{
		if (TileBitmask & MEGALIGHTS_TILE_BITMASK_TEXTURED_RECT_LIGHT)
		{
			Out = InShadingRectTextured;
		}
		else
		{
			Out = InShadingRect;
		}
	}
	return Out;
}

uint GetTileModeFromBitmask(uint TileBitmask)
{
	uint TileMode = TILE_MODE_EMPTY;

	if (TileBitmask & MEGALIGHTS_TILE_BITMASK_COMPLEX_SPECIAL)
	{
		TileMode = BuildTileMode(TileBitmask, TILE_MODE_COMPLEX_SPECIAL_SHADING, TILE_MODE_COMPLEX_SPECIAL_SHADING_RECT, TILE_MODE_COMPLEX_SPECIAL_SHADING_RECT_TEXTURED);
	}
	else if (TileBitmask & MEGALIGHTS_TILE_BITMASK_COMPLEX)
	{
		TileMode = BuildTileMode(TileBitmask, TILE_MODE_COMPLEX_SHADING, TILE_MODE_COMPLEX_SHADING_RECT, TILE_MODE_COMPLEX_SHADING_RECT_TEXTURED);
	}
	else if (TileBitmask & MEGALIGHTS_TILE_BITMASK_SINGLE)
	{
		TileMode = BuildTileMode(TileBitmask, TILE_MODE_SINGLE_SHADING, TILE_MODE_SINGLE_SHADING_RECT, TILE_MODE_SINGLE_SHADING_RECT_TEXTURED);
	}
	else if (TileBitmask & MEGALIGHTS_TILE_BITMASK_SIMPLE)
	{
		TileMode = BuildTileMode(TileBitmask, TILE_MODE_SIMPLE_SHADING, TILE_MODE_SIMPLE_SHADING_RECT, TILE_MODE_SIMPLE_SHADING_RECT_TEXTURED);
	}

	return TileMode;
}

/**
 * Output tile lists for future processing based on the tile bitmask
 */
[numthreads(THREADGROUP_SIZE * THREADGROUP_SIZE, 1, 1)] 
void MegaLightsTileClassificationBuildListsCS(
	uint2 GroupId : SV_GroupID,
	uint GroupThreadId : SV_GroupThreadID)
{
	const uint2 ThreadOffset = ZOrder2D(GroupThreadId, log2(TILE_SIZE));
	uint2 LocalTileCoord = GroupId * TILE_SIZE + ThreadOffset;

	uint TileMode = 0xFF;

#if DOWNSAMPLE_FACTOR_X != 1
	if (all(LocalTileCoord < DownsampledViewSizeInTiles))
	{
		uint MergedTileBitmask = 0;

		for (uint Y = 0; Y < DOWNSAMPLE_FACTOR_Y; ++Y)
		{
			for (uint X = 0; X < DOWNSAMPLE_FACTOR_X; ++X)
			{
				uint2 GatherTileCoord = LocalTileCoord * uint2(DOWNSAMPLE_FACTOR_X, DOWNSAMPLE_FACTOR_Y) + uint2(X, Y) + ViewMinInTiles;
				if (all(GatherTileCoord.xy < ViewMinInTiles + ViewSizeInTiles))
				{
					MergedTileBitmask |= MegaLightsTileBitmask[GatherTileCoord];
				}
			}
		}

		TileMode = GetTileModeFromBitmask(MergedTileBitmask);
	}	
#else
	if (all(LocalTileCoord < ViewSizeInTiles))
	{
		TileMode = GetTileModeFromBitmask(MegaLightsTileBitmask[LocalTileCoord + ViewMinInTiles]);
	}
#endif

	if (TileMode != 0xFF)
	{
		for (uint TileTypeIndex = 0; TileTypeIndex < MAX_USED_TILE_MODE; ++TileTypeIndex)
		{
			uint NumTilesInWave = WaveActiveCountBits(TileTypeIndex == TileMode);
			uint GlobalTileOffset = 0;
			if (WaveIsFirstLane() && NumTilesInWave > 0)
			{
				InterlockedAdd(RWTileAllocator[TileTypeIndex], NumTilesInWave, GlobalTileOffset);
			}

			GlobalTileOffset = WaveReadLaneFirst(GlobalTileOffset);

			if (TileTypeIndex == TileMode)
			{
				RWTileData[OutputTileDataStride * TileTypeIndex + GlobalTileOffset + WavePrefixCountBits(true)] = PackTile(LocalTileCoord);
			}
		}
	}
}

#endif

RWBuffer<uint> RWTileIndirectArgs;
RWBuffer<uint> RWDownsampledTileIndirectArgs;
StructuredBuffer<uint> TileAllocator;
StructuredBuffer<uint> DownsampledTileAllocator;

#ifdef InitTileIndirectArgsCS

[numthreads(THREADGROUP_SIZE, 1, 1)]
void InitTileIndirectArgsCS(uint3 DispatchThreadId : SV_DispatchThreadID)
{
	uint TileMode = DispatchThreadId.x;
	if (TileMode < TILE_MODE_MAX)
	{
		WriteDispatchIndirectArgs(RWTileIndirectArgs, TileMode, TileAllocator[TileMode], 1, 1);
		WriteDispatchIndirectArgs(RWDownsampledTileIndirectArgs, TileMode, DownsampledTileAllocator[TileMode], 1, 1);
	}
}

#endif

#ifdef HairTransmittanceCS

#define VOXEL_TRAVERSAL_DEBUG 0
#define VOXEL_TRAVERSAL_TYPE VOXEL_TRAVERSAL_LINEAR_MIPMAP
#include "../BlueNoise.ush"
#include "MegaLightsRayTracing.ush"
#include "../HairStrands/HairStrandsDeepShadowCommon.ush"
#include "../HairStrands/HairStrandsCommon.ush"
#include "../HairStrands/HairStrandsDeepTransmittanceCommon.ush"
#include "../HairStrands/HairStrandsVisibilityCommon.ush"
#include "../HairStrands/HairStrandsVoxelPageCommon.ush"
#include "../HairStrands/HairStrandsVoxelPageTraversal.ush"
#include "../HairStrands/HairStrandsScreenTracing.ush"

#if INPUT_TYPE != INPUT_TYPE_HAIRSTRANDS
#error INPUT_TYPE is required to be set to INPUT_TYPE_HAIRSTRANDS
#endif

uint bUseScreenTrace;
uint2 NumSamplesPerPixel;
int2 SampleViewMin;
int2 SampleViewSize;
Texture2D<uint> LightSamples;
Texture2D<uint> LightSampleRays;
RWTexture2D<uint> RWTransmittanceMaskTexture;

#define USE_BLUE_NOISE 1
DEFINE_HAIR_BLUE_NOISE_JITTER_FUNCTION

float3 InternalGetHairVoxelJitter(uint2 PixelCoord, uint SampleIndex, uint MaxSampleCount, uint TimeIndex, uint JitterMode)
{
	// Blue noise is cheaper to compute and h
	float3 RandomSample = GetHairBlueNoiseJitter(PixelCoord, SampleIndex, MaxSampleCount, JitterMode > 1 ? 0 : TimeIndex).xyz;
	return JitterMode > 0 ? RandomSample : float3(0,0,0);
}

/**
 * Compute hair transmittance
 */
[numthreads(THREADGROUP_SIZE, THREADGROUP_SIZE, 1)] 
void HairTransmittanceCS(
	uint3 GroupId : SV_GroupID,
	uint3 GroupThreadId : SV_GroupThreadID,
	uint3 DispatchThreadId : SV_DispatchThreadID)
{
	const uint2 SampleCoord = DispatchThreadId.xy + SampleViewMin;
	const uint2 DownsampledScreenCoord = uint2(SampleCoord.x >> NumSamplesPerPixelDivideShift.x, SampleCoord.y >> NumSamplesPerPixelDivideShift.y);
	const uint2 ScreenCoord = DownsampledScreenCoordToScreenCoord(DownsampledScreenCoord);
	const float2 ScreenUV = DownsampledScreenCoordToScreenUV(DownsampledScreenCoord);

	if (all(SampleCoord < SampleViewMin + SampleViewSize))
	{
		const FMegaLightsMaterial Material = LoadMaterial(ScreenUV, ScreenCoord);	
		if (Material.bIsValid)
		{
			const uint2 PixelCoord		= ScreenCoord;
			const float2 UV				= (PixelCoord.xy + float2(0.5f, 0.5f)) / float2(View.BufferSizeAndInvSize.xy);
			const float PixelRadius		= ConvertGivenDepthRadiusForProjectionType(VirtualVoxel.HairCoveragePixelRadiusAtDepth1, Material.Depth);

			const float3 TranslatedWorldPosition = GetTranslatedWorldPositionFromScreenUV(ScreenUV, Material.Depth);

			const FLightSample LightSample = UnpackLightSample(LightSamples[SampleCoord]);
			const FLightSampleRay LightSampleRay = UnpackLightSampleRay(LightSampleRays[SampleCoord]);
			const FLightSampleTrace LightSampleTrace = GetLightSampleTrace(TranslatedWorldPosition, LightSample.LocalLightIndex, LightSampleRay.UV);

			// Early out when the ray is already occluded
			if (LightSampleRay.bCompleted)
			{
				RWTransmittanceMaskTexture[SampleCoord] = 0;					
				return;
			}

			const float3 TranslatedLightPosition= TranslatedWorldPosition + LightSampleTrace.Direction * LightSampleTrace.Distance;
			const float3 LightDirection			= LightSampleTrace.Direction;

			const FVirtualVoxelNodeDesc NodeDesc = UnpackVoxelNode(VirtualVoxel.NodeDescBuffer[Material.MacroGroupId], VirtualVoxel.PageResolution);

			FVirtualVoxelCommonDesc CommonDesc;
			CommonDesc.PageCountResolution	= VirtualVoxel.PageCountResolution;
			CommonDesc.PageTextureResolution= VirtualVoxel.PageTextureResolution;
			CommonDesc.PageResolution		= VirtualVoxel.PageResolution;
			CommonDesc.PageResolutionLog2	= VirtualVoxel.PageResolutionLog2;

			const uint SampleCount = 1; 
			const uint SampleIt = 0;
			float3 SampleRandom = InternalGetHairVoxelJitter(PixelCoord, SampleIt, SampleCount, View.StateFrameIndexMod8, VirtualVoxel.JitterMode);
	
			const float PositionBiasScale = 0.5f;
			const float3 DepthBias = NodeDesc.VoxelWorldSize * (VirtualVoxel.DepthBiasScale_Transmittance * LightDirection + PositionBiasScale*(SampleRandom*2-1));
			const float3 SampleTranslatedWorldPosition = TranslatedWorldPosition + DepthBias;

			FHairTraversalSettings TraversalSettings = InitHairTraversalSettings();
			TraversalSettings.DensityScale = VirtualVoxel.DensityScale_Transmittance;
			TraversalSettings.CountThreshold = GetOpaqueVoxelValue();
			TraversalSettings.DistanceThreshold = 100000;
			TraversalSettings.bDebugEnabled = false;
			TraversalSettings.SteppingScale = VirtualVoxel.SteppingScale_Transmittance;
			TraversalSettings.Random = GetHairVoxelJitter(PixelCoord.xy, View.StateFrameIndexMod8, VirtualVoxel.JitterMode);
			TraversalSettings.TanConeAngle = 0 /*TanLightAngle*/;

			FHairTraversalResult Result = ComputeHairCountVirtualVoxel(
				SampleTranslatedWorldPosition,
				TranslatedLightPosition,
				CommonDesc,
				NodeDesc,
				VirtualVoxel.PageIndexBuffer,
				VirtualVoxel.PageTexture,
				TraversalSettings);

			if (bUseScreenTrace)
			{
				Result.HairCount += ComputeScreenTraceHairCount(PixelCoord, PixelRadius, Material.Depth, TranslatedWorldPosition, LightDirection, SampleIt, SampleCount, VirtualVoxel.JitterMode == 1);
			}

			FHairTransmittanceMask Out;
			Out.HairCount = Result.HairCount;
			Out.Visibility = Result.Visibility;
			RWTransmittanceMaskTexture[SampleCoord] = PackTransmittanceMask(Out);
		}
	}
}

#endif