/*
===========================================================================
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/>.
===========================================================================
*/
// Cinematic Rendering Pipeline - scatter-as-gather depth of field


#include "crp_local.h"
#include "compshaders/crp/fullscreen.h"
#include "compshaders/crp/gatherdof_debug.h"
#include "compshaders/crp/gatherdof_split.h"
#include "compshaders/crp/gatherdof_coc_tile_gen.h"
#include "compshaders/crp/gatherdof_coc_tile_max.h"
#include "compshaders/crp/gatherdof_blur.h"
#include "compshaders/crp/gatherdof_fill.h"
#include "compshaders/crp/gatherdof_combine.h"


#pragma pack(push, 4)

struct DOFDebugRC
{
	uint32_t colorTextureIndex;
	uint32_t depthTextureIndex;
	uint32_t debugMode;
	float focusNearMin;
	float focusNearMax;
	float focusFarMin;
	float focusFarMax;
	float focusDist;
};

struct DOFSplitRC
{
	uint32_t depthTextureIndex;
	uint32_t colorTextureIndex;
	uint32_t nearColorTextureIndex;
	uint32_t farColorTextureIndex;
	uint32_t nearCocTextureIndex;
	uint32_t farCocTextureIndex;
	float focusNearMin;
	float focusNearMax;
	float focusFarMin;
	float focusFarMax;
	float brightnessScale;
};

struct DOFNearCocMaxRC
{
	uint32_t inputTextureIndex;
	uint32_t outputTextureIndex;
	uint32_t samplerIndex;
	int32_t kernelRadius;
	float kernelDirectionX;
	float kernelDirectionY;
};

struct DOFNearCocBlurRC
{
	uint32_t inputTextureIndex;
	uint32_t outputTextureIndex;
	uint32_t samplerIndex;
	int32_t kernelRadius;
	float kernelDirectionX;
	float kernelDirectionY;
};

struct DOFNearCocTileGenRC
{
	uint32_t inputTextureIndex;
	uint32_t outputTextureIndex;
};

struct DOFNearCocTileMaxRC
{
	uint32_t inputTextureIndex;
	uint32_t outputTextureIndex;
	uint32_t samplerIndex; // point/clamp
};

struct DOFBlurRC
{
	uint32_t colorTextureIndex;
	uint32_t nearColorTextureIndex;
	uint32_t nearMaxCocTextureIndex;
	uint32_t nearCocTextureIndex; // blurry
	uint32_t nearOutputTextureIndex;
	uint32_t farColorTextureIndex;
	uint32_t farCocTextureIndex; // sharp
	uint32_t farOutputTextureIndex;
	uint32_t samplerIndex; // linear/clamp
	float brightnessScale;
	float bladeCount;
	float bokehAngleRad;
};

struct DOFFillRC
{
	uint32_t nearInputTextureIndex;
	uint32_t nearOutputTextureIndex;
	uint32_t farInputTextureIndex;
	uint32_t farOutputTextureIndex;
	uint32_t samplerIndex; // point/clamp
};

struct DOFCombineRC
{
	uint32_t nearTextureIndex;
	uint32_t farTextureIndex;
	uint32_t nearCocTextureIndex;
	uint32_t farCocTextureIndex;
	uint32_t sharpTextureIndex;
	uint32_t samplerIndex; // point/clamp
};

#pragma pack(pop)


void GatherDepthOfField::Init()
{
	tileTextureWidth = (uint32_t)(glConfig.vidWidth + 15) / 16;
	tileTextureHeight = (uint32_t)(glConfig.vidHeight + 15) / 16;

	splitPipeline = CreateComputePipeline("DOF split", ShaderByteCode(g_split_cs));
	nearCocTileGenPipeline = CreateComputePipeline("DOF near CoC tile generation", ShaderByteCode(g_coc_tile_gen_cs));
	nearCocTileMaxPipeline = CreateComputePipeline("DOF near CoC tile dilation", ShaderByteCode(g_coc_tile_max_cs));
	blurPipeline = CreateComputePipeline("DOF blur", ShaderByteCode(g_blur_cs));
	fillPipeline = CreateComputePipeline("DOF fill", ShaderByteCode(g_fill_cs));

	{
		GraphicsPipelineDesc desc("DOF combine");
		MakeFullScreenPipeline(desc, ShaderByteCode(g_combine_ps));
		desc.AddRenderTarget(0, crp.renderTargetFormat);
		combinePipeline = CreateGraphicsPipeline(desc);
	}

	{
		GraphicsPipelineDesc desc("DOF viz");
		MakeFullScreenPipeline(desc, ShaderByteCode(g_debug_ps));
		desc.AddRenderTarget(0, crp.renderTargetFormat);
		debugPipeline = CreateGraphicsPipeline(desc);
	}

	{
		TextureDesc desc("DOF far field color", glConfig.vidWidth, glConfig.vidHeight);
		desc.shortLifeTime = true;
		desc.committedResource = true;
		desc.initialState = ResourceStates::UnorderedAccessBit;
		desc.allowedState = ResourceStates::UnorderedAccessBit | ResourceStates::ComputeShaderAccessBit | ResourceStates::PixelShaderAccessBit;
		desc.format = crp.renderTargetFormat;
		farColorTexture = CreateTexture(desc);
		desc.name = "DOF near field color";
		nearColorTexture = CreateTexture(desc);
		desc.name = "DOF near field blurred color";
		nearBlurTexture = CreateTexture(desc);
		desc.name = "DOF far field blurred color";
		farBlurTexture = CreateTexture(desc);

		desc.format = TextureFormat::R8_UNorm;
		desc.name = "DOF near field CoC #1";
		nearCocTexture = CreateTexture(desc);
		desc.name = "DOF near field CoC #2";
		nearCocTexture2 = CreateTexture(desc);
		desc.name = "DOF far field CoC";
		farCocTexture = CreateTexture(desc);

		desc.width = tileTextureWidth;
		desc.height = tileTextureHeight;
		desc.name = "DOF near field CoC tile #1";
		nearCocTileTexture = CreateTexture(desc);
		desc.name = "DOF near field CoC tile #2";
		nearCocTileTexture2 = CreateTexture(desc);
	}
}

void GatherDepthOfField::Draw()
{
	if(crp_dof->integer != DOFMethod::Gather)
	{
		return;
	}

	if(!IsViewportFullscreen(backEnd.viewParms))
	{
		return;
	}

	SCOPED_RENDER_PASS("Gather DOF", 0.125f, 0.125f, 0.25f);

	DrawSplit();
	DrawNearCocTileGen();
	DrawNearCocTileMax();
	DrawBlur();
	DrawFill();
	DrawCombine();
	DrawDebug();
}

void GatherDepthOfField::DrawDebug()
{
	if(crp_dof_overlay->integer == 0)
	{
		return;
	}

	SCOPED_DEBUG_LABEL("DOF Debug", 0.125f, 0.125f, 0.25f);

	crp.SwapRenderTargets();

	CmdBeginBarrier();
	CmdTextureBarrier(crp.GetReadRenderTarget(), ResourceStates::PixelShaderAccessBit);
	CmdTextureBarrier(crp.depthTexture, ResourceStates::PixelShaderAccessBit);
	CmdTextureBarrier(crp.GetWriteRenderTarget(), ResourceStates::RenderTargetBit);
	CmdEndBarrier();

	DOFDebugRC rc = {};
	rc.colorTextureIndex = GetTextureIndexSRV(crp.GetReadRenderTarget());
	rc.depthTextureIndex = GetTextureIndexSRV(crp.depthTexture);
	rc.debugMode = crp_dof_overlay->integer;
	rc.focusNearMin = crp_gatherDof_focusNearDist->value - 0.5f * crp_gatherDof_focusNearRange->value;
	rc.focusNearMax = crp_gatherDof_focusNearDist->value + 0.5f * crp_gatherDof_focusNearRange->value;
	rc.focusFarMin = crp_gatherDof_focusFarDist->value - 0.5f * crp_gatherDof_focusFarRange->value;
	rc.focusFarMax = crp_gatherDof_focusFarDist->value + 0.5f * crp_gatherDof_focusFarRange->value;
	rc.focusDist = 0.5f * (rc.focusNearMax + rc.focusFarMin);

	CmdBindRenderTargets(1, &crp.renderTarget, NULL);
	CmdBindPipeline(debugPipeline);
	CmdSetGraphicsRootConstants(0, sizeof(rc), &rc);
	CmdDraw(3, 0);
}

void GatherDepthOfField::DrawSplit()
{
	SCOPED_DEBUG_LABEL("DOF Split", 0.125f, 0.125f, 0.25f);

	CmdBeginBarrier();
	CmdTextureBarrier(crp.depthTexture, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(crp.renderTarget, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(nearColorTexture, ResourceStates::UnorderedAccessBit);
	CmdTextureBarrier(farColorTexture, ResourceStates::UnorderedAccessBit);
	CmdTextureBarrier(nearCocTexture, ResourceStates::UnorderedAccessBit);
	CmdTextureBarrier(farCocTexture, ResourceStates::UnorderedAccessBit);
	CmdEndBarrier();

	DOFSplitRC rc = {};
	rc.depthTextureIndex = GetTextureIndexSRV(crp.depthTexture);
	rc.colorTextureIndex = GetTextureIndexSRV(crp.renderTarget);
	rc.nearColorTextureIndex = GetTextureIndexUAV(nearColorTexture, 0);
	rc.farColorTextureIndex = GetTextureIndexUAV(farColorTexture, 0);
	rc.nearCocTextureIndex = GetTextureIndexUAV(nearCocTexture, 0);
	rc.farCocTextureIndex = GetTextureIndexUAV(farCocTexture, 0);
	rc.focusNearMin = crp_gatherDof_focusNearDist->value - 0.5f * crp_gatherDof_focusNearRange->value;
	rc.focusNearMax = crp_gatherDof_focusNearDist->value + 0.5f * crp_gatherDof_focusNearRange->value;
	rc.focusFarMin = crp_gatherDof_focusFarDist->value - 0.5f * crp_gatherDof_focusFarRange->value;
	rc.focusFarMax = crp_gatherDof_focusFarDist->value + 0.5f * crp_gatherDof_focusFarRange->value;
	rc.brightnessScale = crp_gatherDof_brightness->value;

	CmdBindPipeline(splitPipeline);
	CmdSetComputeRootConstants(0, sizeof(rc), &rc);
	CmdDispatch((glConfig.vidWidth + 7) / 8, (glConfig.vidHeight + 7) / 8, 1);
}

void GatherDepthOfField::DrawNearCocTileGen()
{
	SCOPED_DEBUG_LABEL("DOF Tile Gen", 0.125f, 0.125f, 0.25f);

	CmdBeginBarrier();
	CmdTextureBarrier(nearCocTexture, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(nearCocTileTexture, ResourceStates::UnorderedAccessBit);
	CmdEndBarrier();

	DOFNearCocTileGenRC rc = {};
	rc.inputTextureIndex = GetTextureIndexSRV(nearCocTexture);
	rc.outputTextureIndex = GetTextureIndexUAV(nearCocTileTexture, 0);

	CmdBindPipeline(nearCocTileGenPipeline);
	CmdSetComputeRootConstants(0, sizeof(rc), &rc);
	CmdDispatch((tileTextureWidth + 7) / 8, (tileTextureHeight + 7) / 8, 1);
}

void GatherDepthOfField::DrawNearCocTileMax()
{
	SCOPED_DEBUG_LABEL("DOF Tile Max", 0.125f, 0.125f, 0.25f);

	CmdBeginBarrier();
	CmdTextureBarrier(nearCocTileTexture, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(nearCocTileTexture2, ResourceStates::UnorderedAccessBit);
	CmdEndBarrier();

	DOFNearCocTileMaxRC rc = {};
	rc.inputTextureIndex = GetTextureIndexSRV(nearCocTileTexture);
	rc.outputTextureIndex = GetTextureIndexUAV(nearCocTileTexture2, 0);
	rc.samplerIndex = GetSamplerIndex(TW_CLAMP_TO_EDGE, TextureFilter::Point);

	CmdBindPipeline(nearCocTileMaxPipeline);
	CmdSetComputeRootConstants(0, sizeof(rc), &rc);
	CmdDispatch((tileTextureWidth + 7) / 8, (tileTextureHeight + 7) / 8, 1);
}

void GatherDepthOfField::DrawBlur()
{
	SCOPED_DEBUG_LABEL("DOF Blur", 0.125f, 0.125f, 0.25f);

	CmdBeginBarrier();
	CmdTextureBarrier(crp.renderTarget, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(nearColorTexture, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(farColorTexture, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(nearCocTexture, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(nearCocTileTexture2, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(farCocTexture, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(nearBlurTexture, ResourceStates::UnorderedAccessBit);
	CmdTextureBarrier(farBlurTexture, ResourceStates::UnorderedAccessBit);
	CmdEndBarrier();

	DOFBlurRC rc = {};
	rc.colorTextureIndex = GetTextureIndexSRV(crp.renderTarget);
	rc.nearColorTextureIndex = GetTextureIndexSRV(nearColorTexture);
	rc.nearMaxCocTextureIndex = GetTextureIndexSRV(nearCocTileTexture2);
	rc.nearCocTextureIndex = GetTextureIndexSRV(nearCocTexture);
	rc.nearOutputTextureIndex = GetTextureIndexUAV(nearBlurTexture, 0);
	rc.farColorTextureIndex = GetTextureIndexSRV(farColorTexture);
	rc.farCocTextureIndex = GetTextureIndexSRV(farCocTexture);
	rc.farOutputTextureIndex = GetTextureIndexUAV(farBlurTexture, 0);
	rc.samplerIndex = GetSamplerIndex(TW_CLAMP_TO_EDGE, TextureFilter::Linear);
	rc.brightnessScale = crp_gatherDof_brightness->value;
	rc.bladeCount = crp_dof_blades->value;
	rc.bokehAngleRad = DEG2RAD(crp_dof_angle->value);

	CmdBindPipeline(blurPipeline);
	CmdSetComputeRootConstants(0, sizeof(rc), &rc);
	CmdDispatch((glConfig.vidWidth + 7) / 8, (glConfig.vidHeight + 7) / 8, 1);
}

void GatherDepthOfField::DrawFill()
{
	SCOPED_DEBUG_LABEL("DOF Fill", 0.125f, 0.125f, 0.25f);

	CmdBeginBarrier();
	CmdTextureBarrier(nearBlurTexture, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(farBlurTexture, ResourceStates::ComputeShaderAccessBit);
	CmdTextureBarrier(nearColorTexture, ResourceStates::UnorderedAccessBit);
	CmdTextureBarrier(farColorTexture, ResourceStates::UnorderedAccessBit);
	CmdEndBarrier();

	DOFFillRC rc = {};
	rc.nearInputTextureIndex = GetTextureIndexSRV(nearBlurTexture);
	rc.farInputTextureIndex = GetTextureIndexSRV(farBlurTexture);
	rc.nearOutputTextureIndex = GetTextureIndexUAV(nearColorTexture, 0);
	rc.farOutputTextureIndex = GetTextureIndexUAV(farColorTexture, 0);
	rc.samplerIndex = GetSamplerIndex(TW_CLAMP_TO_EDGE, TextureFilter::Point);

	CmdBindPipeline(fillPipeline);
	CmdSetComputeRootConstants(0, sizeof(rc), &rc);
	CmdDispatch((glConfig.vidWidth + 7) / 8, (glConfig.vidHeight + 7) / 8, 1);
}

void GatherDepthOfField::DrawCombine()
{
	SCOPED_DEBUG_LABEL("DOF Combine", 0.125f, 0.125f, 0.25f);

	CmdBeginBarrier();
	CmdTextureBarrier(nearColorTexture, ResourceStates::PixelShaderAccessBit);
	CmdTextureBarrier(farColorTexture, ResourceStates::PixelShaderAccessBit);
	CmdTextureBarrier(nearCocTexture, ResourceStates::PixelShaderAccessBit);
	CmdTextureBarrier(farCocTexture, ResourceStates::PixelShaderAccessBit);
	CmdTextureBarrier(crp.GetReadRenderTarget(), ResourceStates::PixelShaderAccessBit);
	CmdTextureBarrier(crp.GetWriteRenderTarget(), ResourceStates::RenderTargetBit);
	CmdEndBarrier();

	DOFCombineRC rc = {};
	rc.nearTextureIndex = GetTextureIndexSRV(nearColorTexture);
	rc.farTextureIndex = GetTextureIndexSRV(farColorTexture);
	rc.nearCocTextureIndex = GetTextureIndexSRV(nearCocTexture);
	rc.farCocTextureIndex = GetTextureIndexSRV(farCocTexture);
	rc.sharpTextureIndex = GetTextureIndexSRV(crp.renderTarget);
	rc.samplerIndex = GetSamplerIndex(TW_CLAMP_TO_EDGE, TextureFilter::Point);

	CmdBindRenderTargets(1, &crp.renderTarget, NULL);
	CmdBindPipeline(combinePipeline);
	CmdSetGraphicsRootConstants(0, sizeof(rc), &rc);
	CmdDraw(3, 0);
}