/*
===========================================================================
Copyright (C) 2023-2024 Gian 'myT' Schellenbaum

This file is part of Challenge Quake 3 (CNQ3).

Challenge Quake 3 is free software; you can redistribute it
and/or modify it under the terms of the GNU General Public License as
published by the Free Software Foundation; either version 2 of the License,
or (at your option) any later version.

Challenge Quake 3 is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with Challenge Quake 3. If not, see <https://www.gnu.org/licenses/>.
===========================================================================
*/
// gather depth of field: near-field and far-field blur


#include "common.hlsli"
#include "gatherdof.hlsli"


cbuffer RootConstants : register(b0)
{
	uint colorTextureIndex;
	uint nearColorTextureIndex;
	uint nearMaxCocTextureIndex; // tile
	uint nearCocTextureIndex; // blurry
	uint nearOutputTextureIndex;
	uint farColorTextureIndex;
	uint farCocTextureIndex; // sharp
	uint farOutputTextureIndex;
	uint samplerIndex; // linear/clamp
	float brightnessScale;
	float bladeCount;
	float bokehAngleRad;
};

// the input is in [0,1]^2, the output polygon is centered at the origin
float2 MapUnitSquareToPolygon(float2 square, float apertureBladeCount, float apertureAngleRad)
{
	// needed to avoid inf/nan propagation through theta for samples
	// that are exactly in the middle of the quad on either axis
	// (i.e. square.x|y == 0.5 gets remapped to 0.0)
	const float epsilon = 0.000001;

	// morph into a square in [-1,1]^2
	square = square * 2.0 - 1.0;

	// morph the square into a disk
	// "A Low Distortion Map Between Disk and Square" by Peter Shirley and Kenneth Chiu
	float radius, angle;
	float2 square2 = square * square;
	if(square2.x > square2.y)
	{
		// left and right quadrants
		radius = square.x;
		angle = (square.y * PI_D4) / (square.x + epsilon);
	}
	else
	{
		// top and bottom quadrants
		radius = square.y;
		angle = PI_D2 - (square.x * PI_D4) / (square.y + epsilon);
	}
	if(radius < 0.0)
	{
		radius = -radius;
		angle += PI;
	}

	// morph the disk into a polygon
	// "Graphics Gems from CryENGINE 3" by Tiago Sousa
	float edgeCount = apertureBladeCount;
	if(edgeCount >= 3.0)
	{
		float num = cos(PI / edgeCount);
		float den0 = PI_M2 / edgeCount;
		float den1 = (angle * edgeCount + PI) / PI_M2;
		float den = angle - (den0 * floor(den1));
		radius *= num / cos(den);
		angle += apertureAngleRad;
	}

	float2 disk;
	sincos(angle, disk.y, disk.x);
	disk *= radius;

	return disk;
}

float4 BlurFarField(Texture2D inTexture, SamplerState samplerState, float coc, float2 tc01, float2 pixelSize)
{
	const int TAP_COUNT_BLUR = 16;
	float2 tcScale = 16.0 * coc * pixelSize;

	float4 result = inTexture.SampleLevel(samplerState, tc01, 0);
	for(int y = 0; y < TAP_COUNT_BLUR; ++y)
	{
		for(int x = 0; x < TAP_COUNT_BLUR; ++x)
		{
			float2 tcQuad = float2(x, y) / float(TAP_COUNT_BLUR - 1);
			float2 tcOffset = MapUnitSquareToPolygon(tcQuad, bladeCount, bokehAngleRad) * tcScale;
			float4 sampleValue = inTexture.SampleLevel(samplerState, tc01 + tcOffset, 0);
			result += sampleValue;
		}
	}

	result /= result.a;

	return result;
}

float4 BlurNearField(Texture2D inTexture, SamplerState samplerState, float tileMaxCoc, float2 tc01, float2 pixelSize)
{
	const int TAP_COUNT_BLUR = 15; // must be odd so we generate 1 sample at 0.5, 0.5 in the quad
	float2 tcScale = 16.0 * tileMaxCoc * pixelSize;
	float insideCount = 0.0;
	float totalCount = 1.0 + float(TAP_COUNT_BLUR * TAP_COUNT_BLUR);
	float4 result = float4(0, 0, 0, 0);
	float weightSum = 0.0;

	for(int y = 0; y < TAP_COUNT_BLUR; ++y)
	{
		for(int x = 0; x < TAP_COUNT_BLUR; ++x)
		{
			float2 tcQuad = float2(x, y) / float(TAP_COUNT_BLUR - 1);
			float2 tcOffset = MapUnitSquareToPolygon(tcQuad, bladeCount, bokehAngleRad) * tcScale;
			float4 sampleValue = inTexture.SampleLevel(samplerState, tc01 + tcOffset, 0);
			float inside = sampleValue.a > 0.0 ? 1.0 : 0.0;
			float brightnessWeight = 1.0 + brightnessScale * Brightness(sampleValue.rgb);
			float colorWeight = (sampleValue.a / tileMaxCoc) * brightnessWeight;
			insideCount += inside;
			weightSum += inside * colorWeight;
			result += inside * float4(colorWeight.xxx, 1) * sampleValue;
		}
	}

	if(insideCount >= 1.0)
	{
		result.rgb /= weightSum;
		result.a /= insideCount;
		result.a *= EaseInOutCubic(saturate(2.0 * (insideCount / totalCount)));
	}
	else
	{
		result = float4(1, 1, 0, 0);
	}

	return result;
}

[numthreads(8, 8, 1)]
void cs(uint3 dtid : SV_DispatchThreadID)
{
	uint2 tc = dtid.xy;
	RWTexture2D<float4> nearOutputTexture = ResourceDescriptorHeap[nearOutputTextureIndex];
	RWTexture2D<float4> farOutputTexture = ResourceDescriptorHeap[farOutputTextureIndex];
	uint width, height;
	farOutputTexture.GetDimensions(width, height);
	if(any(dtid.xy >= uint2(width, height)))
	{
		return;
	}

	SamplerState samplerState = SamplerDescriptorHeap[samplerIndex];
	Texture2D colorTexture = ResourceDescriptorHeap[colorTextureIndex];
	Texture2D nearColorTexture = ResourceDescriptorHeap[nearColorTextureIndex];
	Texture2D farColorTexture = ResourceDescriptorHeap[farColorTextureIndex];
	Texture2D<float> nearCocTexture = ResourceDescriptorHeap[nearCocTextureIndex];
	Texture2D<float> nearMaxCocTexture = ResourceDescriptorHeap[nearMaxCocTextureIndex];
	Texture2D<float> farCocTexture = ResourceDescriptorHeap[farCocTextureIndex];
	RWTexture2D<float4> nearOutTexture = ResourceDescriptorHeap[nearOutputTextureIndex];
	RWTexture2D<float4> farOutTexture = ResourceDescriptorHeap[farOutputTextureIndex];
	float2 tc01 = (float2(dtid.xy) + float2(0.5, 0.5)) / float2(width, height);
	float2 pixelSize = float2(1, 1) / float2(width, height);
	float nearCoc = nearCocTexture.SampleLevel(samplerState, tc01, 0);
	float nearMaxCoc = nearMaxCocTexture.SampleLevel(samplerState, tc01, 0);
	float farCoc = farCocTexture.SampleLevel(samplerState, tc01, 0);
	float4 color = colorTexture.SampleLevel(samplerState, tc01, 0);

	if(nearMaxCoc > 0.0)
	{
		nearOutTexture[tc] = BlurNearField(nearColorTexture, samplerState, nearMaxCoc, tc01, pixelSize);
	}
	else
	{
		// A must be 0 to disable the near field from being blended
		nearOutTexture[tc] = float4(color.rgb, 0);
	}

	if(farCoc > 0.0)
	{
		farOutTexture[tc] = BlurFarField(farColorTexture, samplerState, farCoc, tc01, pixelSize);
	}
	else
	{
		// RGB must be 0 to not mess up the fill pass of neighbor pixels that are inside the near field
		// A must be 0 to disable the far field from being blended
		farOutTexture[tc] = float4(0, 0, 0, 0);
	}
}