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

#pragma once

#ifdef NUM_SAMPLES_PER_VOXEL_1D
	#define NUM_SAMPLES_1D NUM_SAMPLES_PER_VOXEL_1D
#endif

#ifdef NUM_SAMPLES_PER_PIXEL_1D
	#define NUM_SAMPLES_1D NUM_SAMPLES_PER_PIXEL_1D
#endif

struct FCandidateLightSample
{
	uint LocalLightIndex;
	bool bLightWasVisible;
	float Weight;
};

FCandidateLightSample InitCandidateLightSample()
{
	FCandidateLightSample LightSample;
	LightSample.LocalLightIndex = MAX_LOCAL_LIGHT_INDEX;
	LightSample.bLightWasVisible = true;
	LightSample.Weight = 0.0f;
	return LightSample;
}

uint PackCandidateLightSample(FCandidateLightSample LightSample)
{
	uint PackedSample = LightSample.LocalLightIndex & 0xFFFF;
	// Weights are always positive, so reuse the sign bit
	PackedSample |= (f32tof16(LightSample.Weight) & 0x7FFF) << 16;
	PackedSample |= LightSample.bLightWasVisible ? (1U << 31) : 0U;
	return PackedSample;
}

FCandidateLightSample UnpackCandidateLightSample(uint PackedSample)
{
	FCandidateLightSample LightSample;
	LightSample.LocalLightIndex = PackedSample & 0xFFFF;
	LightSample.bLightWasVisible = (PackedSample & (1U << 31)) != 0;
	LightSample.Weight = f16tof32((PackedSample >> 16) & 0x7FFF);
	return LightSample;
}

struct FLightSampler
{
	uint PackedSamples[NUM_SAMPLES_1D];
	float LightIndexRandom[NUM_SAMPLES_1D];
	float WeightSum;
};

FLightSampler InitLightSamplerStratified(float RandomScalar)
{
	FLightSampler LightSampler;
	LightSampler.WeightSum = 0.0f;

	for (uint LightSampleIndex = 0; LightSampleIndex < NUM_SAMPLES_1D; ++LightSampleIndex)
	{
		LightSampler.PackedSamples[LightSampleIndex] = PackCandidateLightSample(InitCandidateLightSample());
		LightSampler.LightIndexRandom[LightSampleIndex] = (RandomScalar + LightSampleIndex) / NUM_SAMPLES_1D;
	}

	return LightSampler;
}

void InitLightSamplerFromSequence(inout FLightSampler LightSampler, uint2 ScreenCoord, uint StateFrameIndex)
{
	LightSampler.WeightSum = 0.0f;

	for (uint LightSampleIndex = 0; LightSampleIndex < NUM_SAMPLES_1D; ++LightSampleIndex)
	{
		FRandomSequence RandSequence = RandomSequenceCreate(uint3(ScreenCoord, StateFrameIndex), LightSampleIndex, NUM_SAMPLES_1D);

		LightSampler.PackedSamples[LightSampleIndex] = PackCandidateLightSample(InitCandidateLightSample());
		LightSampler.LightIndexRandom[LightSampleIndex] = RandSequence.Get1D();
	}
}

float DirectionalLightSampleRatio;
float MinSampleClampingWeight;

void AddLightSample(inout FLightSampler LightSampler, float SampleWeight, uint ForwardLightIndex, bool bWasVisibleInLastFrame, bool bRadialLight)
{
	// Make sure that directional lights don't completely take over the entire sample budget
	if (!bRadialLight && DirectionalLightSampleRatio > 0.0f)
	{
		SampleWeight = min(SampleWeight, max(LightSampler.WeightSum, MinSampleClampingWeight) * DirectionalLightSampleRatio);
	}

	float Tau = LightSampler.WeightSum / (LightSampler.WeightSum + SampleWeight);
	LightSampler.WeightSum += SampleWeight;

	for (uint LightSampleIndex = 0; LightSampleIndex < NUM_SAMPLES_1D; ++LightSampleIndex)
	{
		if (LightSampler.LightIndexRandom[LightSampleIndex] < Tau)
		{
			LightSampler.LightIndexRandom[LightSampleIndex] /= Tau;
		}
		else
		{
			// Select this sample
			LightSampler.LightIndexRandom[LightSampleIndex] = (LightSampler.LightIndexRandom[LightSampleIndex] - Tau) / (1.0f - Tau);

			FCandidateLightSample LightSample = InitCandidateLightSample();
			LightSample.LocalLightIndex = ForwardLightIndex;
			LightSample.bLightWasVisible = bWasVisibleInLastFrame;
			LightSample.Weight = SampleWeight;
			LightSampler.PackedSamples[LightSampleIndex] = PackCandidateLightSample(LightSample);
		}

		LightSampler.LightIndexRandom[LightSampleIndex] = clamp(LightSampler.LightIndexRandom[LightSampleIndex], 0, 0.9999f);
	}
}