// Copyright (C) 2009-present, Panagiotis Christopoulos Charitos and contributors.
// All rights reserved.
// Code licensed under the BSD License.
// http://www.anki3d.org/LICENSE

#pragma anki technique Downscale vert pixel comp
#pragma anki technique Exposure vert pixel comp
#pragma anki technique Upscale vert pixel comp

#include <AnKi/Shaders/QuadVert.hlsl>

// ===========================================================================
// Downscale                                                                 =
// ===========================================================================
#if(ANKI_COMPUTE_SHADER || ANKI_PIXEL_SHADER) && ANKI_TECHNIQUE_Downscale
#	include <AnKi/Shaders/Functions.hlsl>

SamplerState g_linearAnyClampSampler : register(s0);
Texture2D<Vec3> g_tex : register(t0);

struct Constants
{
	Vec2 m_fbSize;
	UVec2 m_padding;
};
ANKI_FAST_CONSTANTS(Constants, g_consts)

#	if ANKI_COMPUTE_SHADER
RWTexture2D<Vec4> g_storageTex : register(u1);
#	endif

#	if ANKI_COMPUTE_SHADER
[numthreads(8, 8, 1)] void main(UVec2 svDispatchThreadId : SV_DISPATCHTHREADID)
{
	const Vec2 uv = (Vec2(svDispatchThreadId) + 0.5) / g_consts.m_fbSize;
#	else
Vec3 main(VertOut input) : SV_TARGET0
{
	const Vec2 uv = input.m_uv;
#	endif

	Vec3 output;
	const F32 weight = 1.0 / 5.0;
	output = g_tex.SampleLevel(g_linearAnyClampSampler, uv, 0.0) * weight;
	output += g_tex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(+1, +1)) * weight;
	output += g_tex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(-1, -1)) * weight;
	output += g_tex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(+1, -1)) * weight;
	output += g_tex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(-1, +1)) * weight;

#	if ANKI_COMPUTE_SHADER
	g_storageTex[svDispatchThreadId] = Vec4(output, 1.0);
#	else
	return output;
#	endif
}
#endif // (ANKI_COMPUTE_SHADER || ANKI_PIXEL_SHADER) && ANKI_TECHNIQUE_Downscale

// ===========================================================================
// Exposure                                                                  =
// ===========================================================================
#if(ANKI_COMPUTE_SHADER || ANKI_PIXEL_SHADER) && ANKI_TECHNIQUE_Exposure
#	include <AnKi/Shaders/TonemappingFunctions.hlsl>
#	include <AnKi/Shaders/Functions.hlsl>

SamplerState g_linearAnyClampSampler : register(s0);
Texture2D<Vec4> g_inTex : register(t0);

#	define TONEMAPPING_REGISTER u0
#	include <AnKi/Shaders/TonemappingResources.hlsl>

struct Consts
{
	F32 m_threshold;
	F32 m_scale;
	F32 m_padding0;
	F32 m_padding1;
};
ANKI_FAST_CONSTANTS(Consts, g_consts)

#	if ANKI_COMPUTE_SHADER
RWTexture2D<Vec4> g_storageTex : register(u1);
#	endif

#	if ANKI_COMPUTE_SHADER
[numthreads(8, 8, 1)] void main(UVec2 svDispatchThreadId : SV_DISPATCHTHREADID)
#	else
Vec3 main(VertOut input) : SV_TARGET0
#	endif
{
#	if ANKI_COMPUTE_SHADER
	Vec2 imgSize;
	g_storageTex.GetDimensions(imgSize.x, imgSize.y);

	const Vec2 uv = (Vec2(svDispatchThreadId) + 0.5) / imgSize;
#	else
	const Vec2 uv = input.m_uv;
#	endif

	const F32 weight = 1.0 / 5.0;
	Vec3 color = g_inTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).rgb * weight;
	color += g_inTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(+1, +1)).rgb * weight;
	color += g_inTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(-1, -1)).rgb * weight;
	color += g_inTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(-1, +1)).rgb * weight;
	color += g_inTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(+1, -1)).rgb * weight;

	color = tonemap<F32>(color, readExposureAndAverageLuminance<F32>().y, g_consts.m_threshold) * g_consts.m_scale;

#	if ANKI_COMPUTE_SHADER
	g_storageTex[svDispatchThreadId] = Vec4(color, 0.0);
#	else
	return color;
#	endif
}
#endif // (ANKI_COMPUTE_SHADER || ANKI_PIXEL_SHADER) && ANKI_TECHNIQUE_Exposure

// ===========================================================================
// Upscale                                                                   =
// ===========================================================================
#if(ANKI_PIXEL_SHADER || ANKI_COMPUTE_SHADER) && ANKI_TECHNIQUE_Upscale
#	include <AnKi/Shaders/Functions.hlsl>

// Constants
#	define ENABLE_CHROMATIC_DISTORTION 1
#	define ENABLE_HALO 1
constexpr U32 kMaxGhosts = 4u;
constexpr F32 kGhostDispersal = 0.7;
constexpr F32 kHaloWidth = 0.4;
constexpr F32 kChromaticDistortion = 3.0;
constexpr F32 kHaloOpacity = 0.5;

SamplerState g_linearAnyClampSampler : register(s0);
Texture2D<Vec4> g_inputTex : register(t0);
Texture2D<Vec3> g_lensDirtTex : register(t1);

#	if ANKI_COMPUTE_SHADER
RWTexture2D<Vec4> g_storageTex : register(u0);
#	endif

Vec3 textureDistorted(Texture2D<Vec4> tex, SamplerState sampl, Vec2 uv,
					  Vec2 direction, // direction of distortion
					  Vec3 distortion) // per-channel distortion factor
{
#	if ENABLE_CHROMATIC_DISTORTION
	return Vec3(tex.SampleLevel(sampl, uv + direction * distortion.r, 0.0).r, tex.SampleLevel(sampl, uv + direction * distortion.g, 0.0).g,
				tex.SampleLevel(sampl, uv + direction * distortion.b, 0.0).b);
#	else
	return tex.SampleLevel(uv, 0.0).rgb;
#	endif
}

Vec3 ssLensFlare(Vec2 uv)
{
	Vec2 textureSize;
	g_inputTex.GetDimensions(textureSize.x, textureSize.y);

	const Vec2 texelSize = 1.0 / textureSize;
	const Vec3 distortion = Vec3(-texelSize.x * kChromaticDistortion, 0.0, texelSize.x * kChromaticDistortion);
	const F32 lensOfHalf = length(Vec2(0.5, 0.5));

	const Vec2 flipUv = Vec2(1.0, 1.0) - uv;

	const Vec2 ghostVec = (Vec2(0.5, 0.5) - flipUv) * kGhostDispersal;

	const Vec2 direction = normalize(ghostVec);
	Vec3 result = Vec3(0.0, 0.0, 0.0);

	// Sample ghosts
	[unroll] for(U32 i = 0u; i < kMaxGhosts; ++i)
	{
		const Vec2 offset = frac(flipUv + ghostVec * F32(i));

		F32 weight = length(Vec2(0.5, 0.5) - offset) / lensOfHalf;
		weight = pow(1.0 - weight, 10.0);

		result += textureDistorted(g_inputTex, g_linearAnyClampSampler, offset, direction, distortion) * weight;
	}

	// Sample halo
#	if ENABLE_HALO
	const Vec2 haloVec = normalize(ghostVec) * kHaloWidth;
	F32 weight = length(Vec2(0.5, 0.5) - frac(flipUv + haloVec)) / lensOfHalf;
	weight = pow(1.0 - weight, 20.0);
	result += textureDistorted(g_inputTex, g_linearAnyClampSampler, flipUv + haloVec, direction, distortion) * (weight * kHaloOpacity);
#	endif

	// Lens dirt
	result *= g_lensDirtTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).rgb;

	return result;
}

Vec3 upscale(Vec2 uv)
{
	const F32 weight = 1.0 / 5.0;
	Vec3 result = g_inputTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).rgb * weight;
	result += g_inputTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(+1, +1)).rgb * weight;
	result += g_inputTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(+1, -1)).rgb * weight;
	result += g_inputTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(-1, -1)).rgb * weight;
	result += g_inputTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0, IVec2(-1, +1)).rgb * weight;

	return result;
}

#	if ANKI_COMPUTE_SHADER
[numthreads(8, 8, 1)] void main(UVec2 svDispatchThreadId : SV_DISPATCHTHREADID)
#	else
Vec3 main(VertOut input) : SV_TARGET0
#	endif
{
#	if ANKI_COMPUTE_SHADER
	Vec2 storageTexSize;
	g_storageTex.GetDimensions(storageTexSize.x, storageTexSize.y);

	const Vec2 uv = (Vec2(svDispatchThreadId) + 0.5) / storageTexSize;
#	else
	const Vec2 uv = input.m_uv;
#	endif

	const Vec3 outColor = ssLensFlare(uv) + upscale(uv);

#	if ANKI_COMPUTE_SHADER
	g_storageTex[svDispatchThreadId] = Vec4(outColor, 0.0);
#	else
	return outColor;
#	endif
}
#endif // (ANKI_PIXEL_SHADER || ANKI_COMPUTE_SHADER) && ANKI_TECHNIQUE_Upscale