/*
===========================================================================
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/>.
===========================================================================
*/
// shared structure for a given scene view


#pragma once


#include "typedefs.h.hlsli"
#if !defined(__cplusplus)
#	include "common.hlsli"
#endif

#if defined(__cplusplus)
#	pragma pack(push, 4)
#endif

// @TODO: move out
struct ParticleEmitter
{
	uint firstIndex;
	uint liveCount; // post-emission, pre-simulation
	uint liveCount2; // post-simulation
	uint deadCount;
	uint totalCount;
	uint emitCount; // how many added this frame
	float maxSeconds;
};

// @TODO: move out
struct Particle
{
	// needed for injection, set by user
	float3 position;
	float radius;
	float3 scattering; // or emissive
	float absorption;
	float anisotropy;
	uint isEmissive;
	// needed for injection, private
	int3 froxelMin;
	int3 froxelMax;
	// extra data for simulation only
	float3 velocity;
	float lifeTime; // in seconds
};

// @TODO: move out
#define MAX_PARTICLES (1 << 20)
#define MAX_PARTICLE_EMITTERS (1 << 10)

// @TODO: move out
struct DynamicLight
{
	float3 position;
	float radius;
	float3 color;
	float padding;
};

#if !defined(__cplusplus)
struct ExtinctionCascade
{
	// copied over
	Texture3D<float> textures[4];
	SamplerState linearClampSampler;
	float4 worldScale;
	float3 cameraPosition;
	float3 textureSize;
	// set by SetSamplingVolume
	uint lowestMipLevel;
	float mipLerp;

	void SetSamplingVolume(float voxelSize)
	{
		float mipLevel = max(log2(voxelSize / worldScale.x), 0.0);
		if(mipLevel <= 2.0)
		{
			lowestMipLevel = uint(floor(mipLevel));
			mipLerp = frac(mipLevel);
		}
		else if(voxelSize >= worldScale.w)
		{
			lowestMipLevel = 3;
			mipLerp = 0.0;
		}
		else
		{
			float base = worldScale.w / worldScale.z;
			float mipLevelFrac = log2(voxelSize / worldScale.z) / log2(base);
			lowestMipLevel = 2;
			mipLerp = mipLevelFrac;
		}
	}

	float ExtinctionAt(float3 position)
	{
		float ext;
		if(ExtinctionAtMip(position, 0, ext))
		{
			return ext;
		}
		if(ExtinctionAtMip(position, 1, ext))
		{
			return ext;
		}
		if(ExtinctionAtMip(position, 2, ext))
		{
			return ext;
		}

		float3 tc = AABoxWorldSpaceToTC(position, cameraPosition, textureSize, worldScale.w);
		ext = textures[3].SampleLevel(linearClampSampler, tc, 0);
		return ext;
	}

	bool ExtinctionAtMip(float3 position, uint mip, out float ext)
	{
		ext = 0.0;
		if(lowestMipLevel == mip)
		{
			float3 tc0 = AABoxWorldSpaceToTC(position, cameraPosition, textureSize, worldScale[mip + 0]);
			if(Is01(tc0)) // @TODO: make sure it's at least a half-texel inside
			{
				float3 tc1 = AABoxWorldSpaceToTC(position, cameraPosition, textureSize, worldScale[mip + 1]);
				float ext0 = textures[mip + 0].SampleLevel(linearClampSampler, tc0, 0);
				float ext1 = textures[mip + 1].SampleLevel(linearClampSampler, tc1, 0);
				ext = lerp(ext0, ext1, mipLerp);
				return true;
			}
		}

		return false;
	}
};
#endif

#if !defined(__cplusplus)
struct SunVShadowCascade
{
	// copied over
	Texture3D<float> textures[4];
	SamplerState linearClampSampler;
	float4 worldScale;
	float3 cameraPosition;
	float3 textureSize;
	float3x3 zToSunMatrix;
	// set by SetSamplingVolume
	uint lowestMipLevel;
	float mipLerp;

	void SetSamplingVolume(float voxelSize)
	{
		float mipLevel = max(log2(voxelSize / worldScale.x), 0.0);
		if(mipLevel <= 2.0)
		{
			lowestMipLevel = uint(floor(mipLevel));
			mipLerp = frac(mipLevel);
		}
		else if(voxelSize >= worldScale.w)
		{
			lowestMipLevel = 3;
			mipLerp = 0.0;
		}
		else
		{
			float base = worldScale.w / worldScale.z;
			float mipLevelFrac = log2(voxelSize / worldScale.z) / log2(base);
			lowestMipLevel = 2;
			mipLerp = mipLevelFrac;
		}
	}

	float TransmittanceAt(float3 positionWS)
	{
		float3 positionSS = cameraPosition + mul(zToSunMatrix, positionWS - cameraPosition);

		float ext;
		if(TransmittanceAtMip(positionSS, 0, ext))
		{
			return ext;
		}
		if(TransmittanceAtMip(positionSS, 1, ext))
		{
			return ext;
		}
		if(TransmittanceAtMip(positionSS, 2, ext))
		{
			return ext;
		}

		float3 tc = AABoxWorldSpaceToTC(positionSS, cameraPosition, textureSize, worldScale.w);
		ext = textures[3].SampleLevel(linearClampSampler, tc, 0);
		return ext;
	}

	bool TransmittanceAtMip(float3 position, uint mip, out float ext)
	{
		ext = 0.0;
		if(lowestMipLevel == mip)
		{
			float3 tc0 = AABoxWorldSpaceToTC(position, cameraPosition, textureSize, worldScale[mip + 0]);
			if(Is01(tc0)) // @TODO: make sure it's at least a half-texel inside
			{
				float3 tc1 = AABoxWorldSpaceToTC(position, cameraPosition, textureSize, worldScale[mip + 1]);
				float ext0 = textures[mip + 0].SampleLevel(linearClampSampler, tc0, 0);
				float ext1 = textures[mip + 1].SampleLevel(linearClampSampler, tc1, 0);
				ext = lerp(ext0, ext1, mipLerp);
				return true;
			}
		}

		return false;
	}
};
#endif

struct SceneView
{
	matrix projectionMatrix;
	matrix invProjectionMatrix;
	matrix viewMatrix;
	matrix invViewMatrix;
	matrix prevViewProjMatrix;
	matrix prevViewMatrix;
	matrix prevProjectionMatrix;
	float3x3 zToSunMatrix;
	float3x3 sunToZMatrix;
	float4 clipPlane;
	float4 debug;
	float3 cameraPosition;
	float sunIntensityDL;
	float3 sunDirection;
	float zNear;
	float3 sunColor;
	float zFar;
	float3 prevCameraPosition;
	float prevZNear;
	float3 cameraForward;
	float prevZFar;
	float3 cameraLeft;
	float sunIntensityVL;
	float3 cameraUp;
	float pointLightIntensityVL;
	float3 linearDepthConstants;
	float frameSeed;
	float3 ambientColor;
	float ambientIntensity;
	float4 extinctionWorldScale;
	uint4 extinctionTextureIndices;
	float4 sunVShadowWorldScale;
	uint4 sunVShadowTextureIndices;
	uint sceneViewIndex;
	uint frameIndex;
	uint depthTextureIndex;
	uint depthMinMaxTextureIndex;
	uint normalTextureIndex;
	uint shadingPositionTextureIndex;
	uint motionVectorTextureIndex;
	uint motionVectorMBTextureIndex;
	uint lightTextureIndex;
	uint sunlightTextureIndex;
	uint tlasBufferIndex;
	uint tlasInstanceBufferIndex;
	uint linearClampSamplerIndex;

#if !defined(__cplusplus)
	float LinearDepth(float zwDepth)
	{
		return ::LinearDepth(zwDepth, linearDepthConstants);
	}

	float3 CamerayRay(float2 ndc)
	{
		float4 pointNDC = float4(ndc, 0.5, 1);
		float4 pointWSw = mul(invViewMatrix, mul(invProjectionMatrix, pointNDC));
		float3 pointWS = pointWSw.xyz / pointWSw.w;
		float3 dir = pointWS - cameraPosition;
		float3 ray = normalize(dir);

		return ray;
	}

#if 0 // exponential depth distribution like The Last of Us Part II

	float TLOU2SliceToViewDepth(float slice, float C)
	{
		const float Q = 1.0;
		float viewDepth = exp2((slice + Q * C) / C) - exp2(Q);

		return viewDepth;
	}

	float TLOU2ViewDepthToSlice(float viewDepth, float C)
	{
		const float Q = 1.0;
		float logArg = exp2(Q) + viewDepth;
		float slice = logArg <= 0.0 ? -666.0 : (C * (log2(logArg) - Q));

		return slice;
	}

	float TLOU2CFromSliceAndViewDepth(float slice, float viewDepth)
	{
		const float Q = 1.0;
		float C = slice / (log2(viewDepth + exp2(Q)) - Q);

		return C;
	}

	float FroxelViewDepthToZ01Ex(float viewDepth, float sliceCount, float zn, float zf)
	{
		float depthRange = zf - zn;
		float C = TLOU2CFromSliceAndViewDepth(sliceCount - 1.0, depthRange);
		float slice = TLOU2ViewDepthToSlice(viewDepth - zn, C);
		float depth01 = slice / sliceCount;

		return depth01;
	}

	float FroxelZ01ToViewDepth(float depth01, float sliceCount)
	{
		float depthRange = zFar - zNear;
		float C = TLOU2CFromSliceAndViewDepth(sliceCount - 1.0, depthRange);
		float slice = depth01 * sliceCount;
		float viewDepth = zNear + TLOU2SliceToViewDepth(slice, C);

		return viewDepth;
	}

#elif 1 // quadratic depth distribution

	float FroxelViewDepthToZ01Ex(float viewDepth, float sliceCount, float zn, float zf)
	{
		float depth01 = (viewDepth - zn) / (zf - zn);
		depth01 = sqrt(depth01);

		return depth01;
	}

	float FroxelZ01ToViewDepth(float depth01, float sliceCount)
	{
		depth01 *= depth01;
		float viewDepth = zNear + (zFar - zNear) * depth01;

		return viewDepth;
	}

#else // linear depth distribution

	float FroxelViewDepthToZ01Ex(float viewDepth, float sliceCount, float zn, float zf)
	{
		float depth01 = (viewDepth - zn) / (zf - zn);

		return depth01;
	}

	float FroxelZ01ToViewDepth(float depth01, float sliceCount)
	{
		float viewDepth = zNear + (zFar - zNear) * depth01;

		return viewDepth;
	}

#endif

	float FroxelViewDepthToZ01(float viewDepth, float sliceCount)
	{
		float depth01 = FroxelViewDepthToZ01Ex(viewDepth, sliceCount, zNear, zFar);

		return depth01;
	}

	float FroxelViewDepthToZ01PrevFrame(float viewDepth, float sliceCount)
	{
		float depth01 = FroxelViewDepthToZ01Ex(viewDepth, sliceCount, prevZNear, prevZFar);

		return depth01;
	}

	float3 FroxelTCToWorldSpace(float3 tc, float3 textureSize)
	{
		float pointDepth = FroxelZ01ToViewDepth(tc.z, textureSize.z);
		float2 xyNDC = TCToNDC(tc.xy);
		float zNDC = PostProjectionDepth(pointDepth, zNear, zFar);
		float4 pointNDC = float4(xyNDC, zNDC, 1);
		float4 pointWSw = mul(invViewMatrix, mul(invProjectionMatrix, pointNDC));
		float3 pointWS = pointWSw.xyz / pointWSw.w;

		return pointWS;
	}

	float3 FroxelIndexToWorldSpace(int3 index, float3 textureSize)
	{
		float3 tc = (float3(index) + float3(0.5, 0.5, 0.5)) / textureSize;
		float3 pointWS = FroxelTCToWorldSpace(tc, textureSize);

		return pointWS;
	}

	float3 FroxelWorldSpaceToTC(float3 positionWS, float3 textureSize)
	{
		float4 positionVSw = mul(viewMatrix, float4(positionWS, 1));
		float viewDepth = -positionVSw.z / positionVSw.w;
		float z01 = FroxelViewDepthToZ01(viewDepth, textureSize.z);
		float4 positionCSw = mul(projectionMatrix, positionVSw);
		float2 xy01 = NDCToTC(positionCSw.xy / positionCSw.w);
		float3 tc = float3(xy01, z01);

		return tc;
	}

	int3 FroxelWorldSpaceToIndex(float3 positionWS, float3 textureSize)
	{
		float3 tc = FroxelWorldSpaceToTC(positionWS, textureSize);
		int3 index = int3(tc * textureSize);

		return index;
	}

	int2 FroxelSphereZExtents(float3 positionWS, float radius, float3 textureSize)
	{
		float4 positionVSw = mul(viewMatrix, float4(positionWS, 1));
		float viewDepth = -positionVSw.z / positionVSw.w;
		float viewDepthMin = viewDepth - radius;
		float viewDepthMax = viewDepth + radius;
		float zMin01 = FroxelViewDepthToZ01(viewDepthMin, textureSize.z);
		float zMax01 = FroxelViewDepthToZ01(viewDepthMax, textureSize.z);
		int2 extents = int2(zMin01 * textureSize.z, ceil(zMax01 * textureSize.z));

		return extents;
	}

	// @TODO: validate new logic
	float3 FroxelReproject01(int3 currIndex, float3 textureSize)
	{
		float3 currTC = (float3(currIndex) + float3(0.5, 0.5, 0.5)) / textureSize;
		float currDepth = FroxelZ01ToViewDepth(currTC.z, textureSize.z);
		float2 xyNDC = TCToNDC(currTC.xy);
		float zNDC = PostProjectionDepth(currDepth, zNear, zFar);
		float4 currNDC = float4(xyNDC, zNDC, 1);
		float4 positionWSw = mul(invViewMatrix, mul(invProjectionMatrix, currNDC));
		float3 positionWS = positionWSw.xyz / positionWSw.w;
		float4 prevVSw = mul(prevViewMatrix, float4(positionWS, 1));
		float prevDepth = -prevVSw.z / prevVSw.w;
		float z01 = FroxelViewDepthToZ01(prevDepth, textureSize.z);
		float4 prevCSw = mul(prevProjectionMatrix, prevVSw);
		float2 xy01 = NDCToTC(prevCSw.xy / prevCSw.w);
		float3 prevTC = float3(xy01, z01);

		return prevTC;
	}

	float FroxelVolume(uint3 index, float3 textureSize)
	{
		float3 tcBase = (float3(index) + float3(0.5, 0.5, 0.5)) / textureSize;
		float3 halfTexel = float3(0.5, 0.5, 0.5) / textureSize;
		float3 posL = FroxelTCToWorldSpace(tcBase + float3(-halfTexel.x, 0, 0), textureSize);
		float3 posR = FroxelTCToWorldSpace(tcBase + float3(halfTexel.x, 0, 0), textureSize);
		float w = distance(posL, posR);
		float3 posU = FroxelTCToWorldSpace(tcBase + float3(0, halfTexel.y, 0), textureSize);
		float3 posD = FroxelTCToWorldSpace(tcBase + float3(0, -halfTexel.y, 0), textureSize);
		float h = distance(posU, posD);
		float3 posF = FroxelTCToWorldSpace(tcBase + float3(0, 0, halfTexel.z), textureSize);
		float3 posB = FroxelTCToWorldSpace(tcBase + float3(0, 0, -halfTexel.z), textureSize);
		float d = distance(posF, posB);
		float volume = w * h * d;

		return volume;
	}

	float3 FroxelAverageDimensions(uint3 index, float3 textureSize)
	{
		float3 tcBase = (float3(index) + float3(0.5, 0.5, 0.5)) / textureSize;
		float3 halfTexel = float3(0.5, 0.5, 0.5) / textureSize;
		float3 posL = FroxelTCToWorldSpace(tcBase + float3(-halfTexel.x, 0, 0), textureSize);
		float3 posR = FroxelTCToWorldSpace(tcBase + float3(halfTexel.x, 0, 0), textureSize);
		float w = distance(posL, posR);
		float3 posU = FroxelTCToWorldSpace(tcBase + float3(0, halfTexel.y, 0), textureSize);
		float3 posD = FroxelTCToWorldSpace(tcBase + float3(0, -halfTexel.y, 0), textureSize);
		float h = distance(posU, posD);
		float3 posF = FroxelTCToWorldSpace(tcBase + float3(0, 0, -halfTexel.z), textureSize);
		float3 posB = FroxelTCToWorldSpace(tcBase + float3(0, 0, halfTexel.z), textureSize);
		float d = distance(posF, posB);
		float3 dimensions = float3(w, h, d);

		return dimensions;
	}
	
	float3 FroxelMaxDimensions(uint3 index, float3 textureSize)
	{
		float3 tcBase = (float3(index) + float3(0.5, 0.5, 0.5)) / textureSize;
		float3 halfTexel = float3(0.5, 0.5, 0.5) / textureSize;
		float3 posL = FroxelTCToWorldSpace(tcBase + float3(-halfTexel.x, 0, halfTexel.z), textureSize);
		float3 posR = FroxelTCToWorldSpace(tcBase + float3(halfTexel.x, 0, halfTexel.z), textureSize);
		float w = distance(posL, posR);
		float3 posU = FroxelTCToWorldSpace(tcBase + float3(0, halfTexel.y, halfTexel.z), textureSize);
		float3 posD = FroxelTCToWorldSpace(tcBase + float3(0, -halfTexel.y, halfTexel.z), textureSize);
		float h = distance(posU, posD);
		float3 posF = FroxelTCToWorldSpace(tcBase + float3(0, 0, -halfTexel.z), textureSize);
		float3 posB = FroxelTCToWorldSpace(tcBase + float3(0, 0, halfTexel.z), textureSize);
		float d = distance(posF, posB);
		float3 dimensions = float3(w, h, d);

		return dimensions;
	}

#if 0
	void FroxelAABB(out float boxMin, out float3 boxMax, int3 index, float3 textureSize)
	{
		float3 tc = (float3(index) + float3(0.5, 0.5, 0.5)) / textureSize;
		float3 halfTexel = float3(0.5, 0.5, 0.5) / textureSize;
		float3 pointWS;
		ClearBoundingBox(boxMin, boxMax);
		pointWS = FroxelTCToWorldSpace(tc + float3(-halfTexel.x, 0, 0), textureSize);
		ExpandBoundingBox(boxMin, boxMax, pointWS);
		pointWS = FroxelTCToWorldSpace(tc + float3(halfTexel.x, 0, 0), textureSize);
		ExpandBoundingBox(boxMin, boxMax, pointWS);
		pointWS = FroxelTCToWorldSpace(tc + float3(0, -halfTexel.y, 0), textureSize);
		ExpandBoundingBox(boxMin, boxMax, pointWS);
		pointWS = FroxelTCToWorldSpace(tc + float3(0, halfTexel.y, 0), textureSize);
		ExpandBoundingBox(boxMin, boxMax, pointWS);
		pointWS = FroxelTCToWorldSpace(tc + float3(0, 0, -halfTexel.z), textureSize);
		ExpandBoundingBox(boxMin, boxMax, pointWS);
		pointWS = FroxelTCToWorldSpace(tc + float3(0, 0, halfTexel.z), textureSize);
		ExpandBoundingBox(boxMin, boxMax, pointWS);
	}
#endif

	float3 ExtinctionIndexToWorldSpace(int3 index, float3 textureSize, float worldScale)
	{
		return AABoxIndexToWorldSpace(index, cameraPosition, textureSize, worldScale);
	}

	float3 ExtinctionTCToWorldSpace(float3 tc, float3 textureSize, float worldScale)
	{
		return AABoxTCToWorldSpace(tc, cameraPosition, textureSize, worldScale);
	}

	float3 ExtinctionWorldSpaceToTC(float3 position, float3 textureSize, float worldScale)
	{
		return AABoxWorldSpaceToTC(position, cameraPosition, textureSize, worldScale);
	}

	int3 ExtinctionWorldSpaceToIndex(float3 position, float3 textureSize, float worldScale)
	{
		return AABoxWorldSpaceToIndex(position, cameraPosition, textureSize, worldScale);
	}

	ExtinctionCascade GetExtinctionCascade(float voxelSize)
	{
		ExtinctionCascade cascade;
		cascade.textures[0] = ResourceDescriptorHeap[extinctionTextureIndices.x];
		cascade.textures[1] = ResourceDescriptorHeap[extinctionTextureIndices.y];
		cascade.textures[2] = ResourceDescriptorHeap[extinctionTextureIndices.z];
		cascade.textures[3] = ResourceDescriptorHeap[extinctionTextureIndices.w];
		cascade.linearClampSampler = SamplerDescriptorHeap[linearClampSamplerIndex];
		cascade.worldScale = extinctionWorldScale;
		cascade.cameraPosition = cameraPosition;
		cascade.textureSize = GetTextureSize(cascade.textures[0]);
		cascade.SetSamplingVolume(voxelSize);

		return cascade;
	}

	SunVShadowCascade GetSunVShadowCascade(float voxelSize)
	{
		SunVShadowCascade cascade;
		cascade.textures[0] = ResourceDescriptorHeap[sunVShadowTextureIndices.x];
		cascade.textures[1] = ResourceDescriptorHeap[sunVShadowTextureIndices.y];
		cascade.textures[2] = ResourceDescriptorHeap[sunVShadowTextureIndices.z];
		cascade.textures[3] = ResourceDescriptorHeap[sunVShadowTextureIndices.w];
		cascade.linearClampSampler = SamplerDescriptorHeap[linearClampSamplerIndex];
		cascade.worldScale = sunVShadowWorldScale;
		cascade.cameraPosition = cameraPosition;
		cascade.textureSize = GetTextureSize(cascade.textures[0]);
		cascade.zToSunMatrix = zToSunMatrix;
		cascade.SetSamplingVolume(voxelSize);

		return cascade;
	}
#endif
};

#if defined(__cplusplus)
#	pragma pack(pop)
#endif

#if defined(__cplusplus)
	static_assert(sizeof(DynamicLight) == 32, "sizeof(DynamicLight) is wrong");
#endif

#if !defined(__cplusplus)
	SceneView GetSceneView()
	{
		StructuredBuffer<SceneView> sceneViewBuffer = ResourceDescriptorHeap[0];
		SceneView sceneView = sceneViewBuffer[0];

		return sceneView;
	}
#endif