/*
===========================================================================
Copyright (C) 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 - accumulation depth of field


#include "crp_local.h"
#include "compshaders/crp/fullscreen.h"
#include "compshaders/crp/accumdof_accum.h"
#include "compshaders/crp/accumdof_norm.h"
#include "compshaders/crp/accumdof_debug.h"


#pragma pack(push, 4)

struct DOFAccumRC
{
	uint32_t textureIndex;
};

struct DOFNormRC
{
	uint32_t textureIndex;
};

struct DOFDebugRC
{
	float mvp[16]; // displaced view, to project to CS
	float invMvp[16]; // main view, to unproject to WS
	uint32_t colorTextureIndex;
	uint32_t depthTextureIndex;
	uint32_t debugMode; // 1: colorized coc, 2: constant intensity far field
	uint32_t tcScale;
	float focusDist;
	float maxNearCocCS;
	float maxFarCocCS;
};

#pragma pack(pop)


// the input is in [0,1]^2, the output polygon is centered at the origin
static void MapUnitSquareToPolygon(const vec2_t square01, float apertureBladeCount, float apertureAngleRad, vec2_t polygon)
{
	// 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.5f gets remapped to 0.0f)
	const float epsilon = 0.000001f;

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

	// morph the square into a disk
	// "A Low Distortion Map Between Disk and Square" by Peter Shirley and Kenneth Chiu
	float radius, angle;
	vec2_t square2;
	square2[0] = square[0] * square[0];
	square2[1] = square[1] * square[1];
	if(square2[0] > square2[1])
	{
		// left and right quadrants
		radius = square[0];
		angle = (square[1] * M_PI_D4) / (square[0] + epsilon);
	}
	else
	{
		// top and bottom quadrants
		radius = square[1];
		angle = M_PI_D2 - (square[0] * M_PI_D4) / (square[1] + epsilon);
	}
	if(radius < 0.0f)
	{
		radius = -radius;
		angle += M_PI;
	}

	// morph the disk into a polygon
	// "Graphics Gems from CryENGINE 3" by Tiago Sousa
	const float edgeCount = apertureBladeCount;
	if(edgeCount >= 3.0f)
	{
		const float num = cosf(M_PI / edgeCount);
		const float den0 = M_PI_M2 / edgeCount;
		const float den1 = (angle * edgeCount + M_PI) / M_PI_M2;
		const float den = angle - (den0 * floorf(den1));
		radius *= num / cosf(den);
		angle += apertureAngleRad;
	}

	polygon[0] = cosf(angle) * radius;
	polygon[1] = sinf(angle) * radius;
}

static int GetResolutionScale()
{
	switch(crp_accumDof_preview->integer)
	{
		case 0: return 1;
		case 1: return 2;
		case 2: return 4;
		default: Q_assert(0); return 1;
	}
}


void AccumDepthOfField::Init()
{
	{
		GraphicsPipelineDesc desc("DOF Accumulate");
		MakeFullScreenPipeline(desc, ShaderByteCode(g_accum_ps));
		desc.AddRenderTarget(GLS_SRCBLEND_ONE | GLS_DSTBLEND_ONE, crp.renderTargetFormat);
		accumPipeline = CreateGraphicsPipeline(desc);
	}

	{
		GraphicsPipelineDesc desc("DOF Normalize");
		MakeFullScreenPipeline(desc, ShaderByteCode(g_norm_ps));
		desc.AddRenderTarget(0, crp.renderTargetFormat);
		normPipeline = CreateGraphicsPipeline(desc);
	}

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

	{
		TextureDesc desc("DOF accumulation", glConfig.vidWidth, glConfig.vidHeight);
		desc.initialState = ResourceStates::RenderTargetBit;
		desc.allowedState = ResourceStates::RenderTargetBit | ResourceStates::PixelShaderAccessBit;
		Vector4Clear(desc.clearColor);
		desc.usePreferredClearValue = true;
		desc.committedResource = true;
		desc.format = crp.renderTargetFormat;
		desc.shortLifeTime = true;
		accumTexture = CreateTexture(desc);
	}
}

void AccumDepthOfField::Begin(const drawSceneViewCommand_t& cmd)
{
	CmdSetViewportAndScissor(0, 0, glConfig.vidWidth, glConfig.vidHeight);

	CmdBeginBarrier();
	CmdTextureBarrier(accumTexture, ResourceStates::RenderTargetBit);
	CmdEndBarrier();

	CmdClearColorTarget(accumTexture, vec4_zero);

	// project a point a few units in front of the main view onto a viewpoint with maximum displacement
	// the distance to the clip space center (i.e. the 2D origin) is the maximum expected CoC
	vec3_t testPoint;
	VectorMA(cmd.viewParms.world.viewOrigin, 16.0f, cmd.viewParms.orient.axis[0], testPoint);
	drawSceneViewCommand_t newCmd;
	FixCommand(newCmd, cmd, 0, 0);
	vec4_t eye, clip;
	R_TransformModelToClip(testPoint, newCmd.viewParms.world.modelMatrix, newCmd.viewParms.projectionMatrix, eye, clip);
	Vector4Scale(clip, 1.0f / clip[3], clip);
	maxNearCocCS = sqrtf(clip[0] * clip[0] + clip[1] * clip[1]);

	// same thing for a point far away
	VectorMA(cmd.viewParms.world.viewOrigin, 69420.0f, cmd.viewParms.orient.axis[0], testPoint);
	FixCommand(newCmd, cmd, 0, 0);
	R_TransformModelToClip(testPoint, newCmd.viewParms.world.modelMatrix, newCmd.viewParms.projectionMatrix, eye, clip);
	Vector4Scale(clip, 1.0f / clip[3], clip);
	maxFarCocCS = sqrtf(clip[0] * clip[0] + clip[1] * clip[1]);
}

uint32_t AccumDepthOfField::GetSampleCount()
{
	switch(crp_accumDof_preview->integer)
	{
		case 1: return 3;
		case 2: return 5;
		default: break;
	}

	return 2 * crp_accumDof_samples->integer + 1;
}

void AccumDepthOfField::FixCommand(drawSceneViewCommand_t& newCmd, const drawSceneViewCommand_t& cmd, uint32_t x, uint32_t y)
{
	const float radius = crp_accumDof_radius->value;
	const float tcScale = 1.0f / (float)(GetSampleCount() - 1);

	vec2_t square01;
	square01[0] = x * tcScale;
	square01[1] = y * tcScale;
	vec2_t polygon;
	MapUnitSquareToPolygon(square01, crp_dof_blades->value, DEG2RAD(crp_dof_angle->value), polygon);

	// 0=forward, 1=left, 2=up
	vec3_t axis[3];
	VectorCopy(cmd.viewParms.orient.axis[0], axis[0]);
	VectorCopy(cmd.viewParms.orient.axis[1], axis[1]);
	VectorCopy(cmd.viewParms.orient.axis[2], axis[2]);

	vec3_t origin;
	VectorMA(cmd.viewParms.world.viewOrigin, radius * polygon[0], axis[1], origin);
	VectorMA(origin, radius * polygon[1], axis[2], origin);

	vec3_t focusPoint;
	VectorMA(cmd.viewParms.world.viewOrigin, crp_accumDof_focusDist->value, axis[0], focusPoint);
	VectorSubtract(focusPoint, origin, axis[0]); // forward
	VectorNormalize(axis[0]);
	CrossProduct(axis[2], axis[0], axis[1]); // left
	VectorNormalize(axis[1]);
	CrossProduct(axis[0], axis[1], axis[2]); // up
	VectorNormalize(axis[2]);

	newCmd = cmd;
	VectorCopy(origin, newCmd.viewParms.orient.origin);
	VectorCopy(origin, newCmd.viewParms.world.viewOrigin);
	R_CreateWorldModelMatrix(origin, axis, newCmd.viewParms.world.modelMatrix);
	newCmd.viewParms.viewportWidth /= GetResolutionScale();
	newCmd.viewParms.viewportHeight /= GetResolutionScale();

	if(x == 0 && y == 0)
	{
		memcpy(modelViewMatrix, newCmd.viewParms.world.modelMatrix, sizeof(modelViewMatrix));
		memcpy(projMatrix, newCmd.viewParms.projectionMatrix, sizeof(projMatrix));
	}
}

void AccumDepthOfField::Accumulate()
{
	srp.renderMode = RenderMode::None;

	SCOPED_RENDER_PASS("DOF Accum", 0.5f, 1.0f, 0.5f);

	CmdBeginBarrier();
	CmdTextureBarrier(crp.renderTarget, ResourceStates::PixelShaderAccessBit);
	CmdTextureBarrier(accumTexture, ResourceStates::RenderTargetBit);
	CmdEndBarrier();

	DOFAccumRC rc = {};
	rc.textureIndex = GetTextureIndexSRV(crp.renderTarget);

	CmdBindRenderTargets(1, &accumTexture, NULL);
	CmdBindPipeline(accumPipeline);
	CmdSetGraphicsRootConstants(0, sizeof(rc), &rc);
	CmdDraw(3, 0);
}

void AccumDepthOfField::Normalize()
{
	srp.renderMode = RenderMode::None;

	{
		SCOPED_RENDER_PASS("DOF Norm", 0.5f, 1.0f, 0.5f);

		CmdBeginBarrier();
		CmdTextureBarrier(accumTexture, ResourceStates::PixelShaderAccessBit);
		CmdTextureBarrier(crp.renderTarget, ResourceStates::RenderTargetBit);
		CmdEndBarrier();

		DOFNormRC rc = {};
		rc.textureIndex = GetTextureIndexSRV(accumTexture);

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

	if(crp_accumDof_preview->integer)
	{
		crp.SwapRenderTargets();
		const float scale = 1.0f / (float)GetResolutionScale();
		const vec2_t tcScale = { scale, scale };
		crp.Blit(crp.GetWriteRenderTarget(), crp.GetReadRenderTarget(), "DOF Upscale", true, tcScale, vec2_zero);
	}
}

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

	srp.renderMode = RenderMode::None;

	SCOPED_RENDER_PASS("DOF Debug", 0.5f, 1.0f, 0.5f);

	crp.SwapRenderTargets();

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

	float mvp[16];
	DOFDebugRC rc = {};
	rc.colorTextureIndex = GetTextureIndexSRV(crp.GetReadRenderTarget());
	rc.depthTextureIndex = GetTextureIndexSRV(crp.depthTexture);
	rc.debugMode = crp_dof_overlay->integer;
	rc.focusDist = crp_accumDof_focusDist->value;
	rc.maxNearCocCS = maxNearCocCS;
	rc.maxFarCocCS = maxFarCocCS;
	rc.tcScale = GetResolutionScale();
	R_MultMatrix(modelViewMatrix, projMatrix, rc.mvp);
	R_MultMatrix(backEnd.viewParms.world.modelMatrix, backEnd.viewParms.projectionMatrix, mvp);
	R_InvMatrix(mvp, rc.invMvp);

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