// Copyright Epic Games, Inc. All Rights Reserved.

#include "Common.ush"
#include "ScreenPass.ush"
#include "PostProcessCommon.ush"
#include "PaniniProjection.ush"
#include "TextureSampling.ush"
#include "LensDistortion.ush"

#define UPSCALE_METHOD_NEAREST 0
#define UPSCALE_METHOD_BILINEAR 1
#define UPSCALE_METHOD_DIRECTIONAL 2
#define UPSCALE_METHOD_CATMULL_ROM 3
#define UPSCALE_METHOD_LANCZOS 4
#define UPSCALE_METHOD_GAUSSIAN 5
#define UPSCALE_METHOD_SMOOTHSTEP 6

// only for MainPS
float UpscaleSoftness;

SCREEN_PASS_TEXTURE_VIEWPORT(Input)
SCREEN_PASS_TEXTURE_VIEWPORT(Output)

Texture2D<float2> DistortingDisplacementTexture;
SamplerState DistortingDisplacementSampler;

Texture2D SceneColorTexture;
SamplerState SceneColorSampler;

// Point-sampled version (used on mobile). Needs to be a separate texture because of OpenGL fused samplers.
Texture2D PointSceneColorTexture;
SamplerState PointSceneColorSampler;

//in a multiview case, PointSceneColorTexture is an array texture
Texture2DArray PointSceneColorTextureArray;

uint2 GridDimensions;
uint bInvertAlpha;

// Do a smoothstep(x) = 3 x^2 - 2 x^3
float2 GetSmoothstepUV(float2 LinearUV, float2 TextureSize, float2 TextureInvSize)
{
	// Top left cornered pixel coordinate to sample.
	float2 PixelCoord = LinearUV * TextureSize - 0.5;

	// Index of the top left pixel used in the bilinear interpolation.
	float2 TopLeftPixelCoord = floor(PixelCoord);

	// Interpolation factors in the 2x2 quad.
	float2 PixelInterp = PixelCoord - TopLeftPixelCoord;
	
	// New interpolation factors in the 2x2 quad with smoothstep.
	float2 SmoothPixelInterp = PixelInterp * PixelInterp * (3 - 2 * PixelInterp);

	// Returns new UV coordinate.
	return TextureInvSize * (TopLeftPixelCoord + SmoothPixelInterp + 0.5);
}

float Luma(float3 Color) 
{
#if UE_LEGACY_LUMINANCE_FACTORS || WORKING_COLOR_SPACE_IS_SRGB
	// Note: In this case, the previous (legacy) luminance factors correctly used the sRGB definition.
	float3 LuminanceFactors = float3(0.2126390059, 0.7151686788, 0.0721923154);
#else
	float3 LuminanceFactors = float3(WORKING_COLOR_SPACE_RGB_TO_XYZ_MAT._m10_m11_m12);
#endif

	return dot(Color, LuminanceFactors);
}

float Gaussian(float Scale, float2 Offset) 
{
	return exp2(Scale * dot(Offset, Offset));
}

// vertex shader entry point
void MainVS(
	in uint VertexId : SV_VertexID,
	in FStereoVSInput StereoInput,
	out noperspective float4 OutTexCoord : TEXCOORD0,
	out FStereoVSOutput StereoOutput,
	out float4 OutPosition : SV_POSITION)
{
	StereoSetupVS(StereoInput, StereoOutput);
	
	float2 TexCoord = float2(VertexId % (GridDimensions.x + 1), VertexId / (GridDimensions.x + 1)) / float2(GridDimensions);

	// still in 0..1 range
	float4 Position = float4(TexCoord.x, TexCoord.y, 0, 1);

	// distort pos
	Position.xy = ApplyLensDistortionOnViewportUV(DistortingDisplacementTexture, DistortingDisplacementSampler, TexCoord);

	DrawRectangle(Position, TexCoord, OutPosition, OutTexCoord.xy);
	OutTexCoord.zw =  OutPosition.xy;
}

float4 SampleSceneColorRGBA(float2 BufferUV)
{
	BufferUV = clamp(BufferUV, Input_UVViewportBilinearMin, Input_UVViewportBilinearMax);

	return SceneColorTexture.SampleLevel(SceneColorSampler, BufferUV, 0).rgba;
}

float4 AccumulateAndApplyWeight(in float4 DataRGBA, in float Weight, inout float WeightsSum)
{
	WeightsSum += Weight;
	return DataRGBA * Weight;
}

void MainPS(noperspective float4 UVAndScreenPos : TEXCOORD0, FStereoPSInput StereoInput, float4 SvPosition : SV_POSITION, out float4 OutColor : SV_Target0)
{
	StereoSetupPS(StereoInput);
	const uint EyeIndex = GetEyeIndex(StereoInput);
	
	OutColor = 0;

#if METHOD == UPSCALE_METHOD_NEAREST
	// Nearest sampling (not blurry but blocky, more for testing)
	#if ES3_1_PROFILE
		#if MOBILE_MULTI_VIEW
			OutColor = Texture2DArraySample(PointSceneColorTextureArray, PointSceneColorSampler, float3(UVAndScreenPos.xy,EyeIndex));
		#else
			OutColor = Texture2DSample(PointSceneColorTexture, PointSceneColorSampler, UVAndScreenPos.xy);
		#endif
	#else
		#if MOBILE_MULTI_VIEW
			OutColor = PointSceneColorTextureArray.SampleLevel(PointSceneColorSampler, vec3(UVAndScreenPos.xy,EyeIndex), 0, int2(0, 0));
		#else
			OutColor = PointSceneColorTexture.SampleLevel(PointSceneColorSampler, UVAndScreenPos.xy, 0, int2(0, 0));
		#endif
	#endif

#elif METHOD == UPSCALE_METHOD_BILINEAR
	// Bilinear (fast, aliasing)
	OutColor = SampleSceneColorRGBA(UVAndScreenPos.xy);

#elif METHOD == UPSCALE_METHOD_DIRECTIONAL
	// Directional blur with unsharp mask upsample.
	float2 UV = UVAndScreenPos.xy;
	float X = 0.5;
	float4 ColorNW = SampleSceneColorRGBA(UV + float2(-X, -X) * Input_ExtentInverse);
	float4 ColorNE = SampleSceneColorRGBA(UV + float2( X, -X) * Input_ExtentInverse);
	float4 ColorSW = SampleSceneColorRGBA(UV + float2(-X,  X) * Input_ExtentInverse);
	float4 ColorSE = SampleSceneColorRGBA(UV + float2( X,  X) * Input_ExtentInverse);
	OutColor = (ColorNW * 0.25) + (ColorNE * 0.25) + (ColorSW * 0.25) + (ColorSE * 0.25);
	
	float LumaNW = Luma(ColorNW.rgb);
	float LumaNE = Luma(ColorNE.rgb);
	float LumaSW = Luma(ColorSW.rgb);
	float LumaSE = Luma(ColorSE.rgb);

	float2 IsoBrightnessDir;
	float DirSWMinusNE = LumaSW - LumaNE;
	float DirSEMinusNW = LumaSE - LumaNW;
	IsoBrightnessDir.x = DirSWMinusNE + DirSEMinusNW;
	IsoBrightnessDir.y = DirSWMinusNE - DirSEMinusNW;

	// avoid NaN on zero vectors by adding 2^-24 (float ulp when length==1, and also minimum representable half)
	IsoBrightnessDir = IsoBrightnessDir * (0.125 * rsqrt(dot(IsoBrightnessDir, IsoBrightnessDir) + 6e-8));

	float4 ColorN = SampleSceneColorRGBA(UV - IsoBrightnessDir * Input_ExtentInverse);
	float4 ColorP = SampleSceneColorRGBA(UV + IsoBrightnessDir * Input_ExtentInverse);

	float UnsharpMask = 0.25;
	OutColor = (ColorN + ColorP) * ((UnsharpMask + 1.0) * 0.5) - (OutColor * UnsharpMask);

#elif METHOD == UPSCALE_METHOD_CATMULL_ROM
	// Bicubic Catmull-Rom in five samples
	FCatmullRomSamples Samples = GetBicubic2DCatmullRomSamples(UVAndScreenPos.xy, Input_Extent, Input_ExtentInverse);
	for (uint i = 0; i < Samples.Count; i++)
	{
		OutColor += SampleSceneColorRGBA(Samples.UV[i]) * Samples.Weight[i];
	}
	OutColor *= Samples.FinalMultiplier;

#elif METHOD == UPSCALE_METHOD_LANCZOS
	{
		// Lanczos 3
		float2 UV = UVAndScreenPos.xy * Input_Extent;
		float2 tc = floor(UV - 0.5) + 0.5;
		float2 f = UV - tc + 2;

		// compute at f, f-1, f-2, f-3, f-4, and f-5 using trig angle addition
		float2 fpi = f*PI, fpi3 = f * (PI / 3.0);
		float2 sinfpi = sin(fpi), sinfpi3 = sin(fpi3), cosfpi3 = cos(fpi3);
		const float r3 = sqrt(3.0);
		float2 w0 = ( sinfpi *       sinfpi3              ) / ( f       * f       );
		float2 w1 = (-sinfpi * (     sinfpi3 - r3*cosfpi3)) / ((f - 1.0)*(f - 1.0));
		float2 w2 = ( sinfpi * (    -sinfpi3 - r3*cosfpi3)) / ((f - 2.0)*(f - 2.0));
		float2 w3 = (-sinfpi * (-2.0*sinfpi3             )) / ((f - 3.0)*(f - 3.0));
		float2 w4 = ( sinfpi * (    -sinfpi3 + r3*cosfpi3)) / ((f - 4.0)*(f - 4.0));
		float2 w5 = (-sinfpi * (     sinfpi3 + r3*cosfpi3)) / ((f - 5.0)*(f - 5.0));

		// use bilinear texture weights to merge center two samples in each dimension
		float2 Weight[5];
		Weight[0] = w0;
		Weight[1] = w1;
		Weight[2] = w2 + w3;
		Weight[3] = w4;
		Weight[4] = w5;

		float2 Sample[5];
		Sample[0] = Input_ExtentInverse * (tc - 2);
		Sample[1] = Input_ExtentInverse * (tc - 1);
		Sample[2] = Input_ExtentInverse * (tc + w3 / Weight[2]);
		Sample[3] = Input_ExtentInverse * (tc + 2);
		Sample[4] = Input_ExtentInverse * (tc + 3);

		OutColor = 0;
		float WeightsSum = 0;

		// 5x5 footprint with corners dropped to give 13 texture taps
		OutColor += AccumulateAndApplyWeight(SampleSceneColorRGBA(float2(Sample[0].x, Sample[2].y)), Weight[0].x * Weight[2].y, WeightsSum);

		OutColor += AccumulateAndApplyWeight(SampleSceneColorRGBA(float2(Sample[1].x, Sample[1].y)), Weight[1].x * Weight[1].y, WeightsSum);
		OutColor += AccumulateAndApplyWeight(SampleSceneColorRGBA(float2(Sample[1].x, Sample[2].y)), Weight[1].x * Weight[2].y, WeightsSum);
		OutColor += AccumulateAndApplyWeight(SampleSceneColorRGBA(float2(Sample[1].x, Sample[3].y)), Weight[1].x * Weight[3].y, WeightsSum);

		OutColor += AccumulateAndApplyWeight( SampleSceneColorRGBA(float2(Sample[2].x, Sample[0].y)), Weight[2].x * Weight[0].y, WeightsSum);
		OutColor += AccumulateAndApplyWeight( SampleSceneColorRGBA(float2(Sample[2].x, Sample[1].y)), Weight[2].x * Weight[1].y, WeightsSum);
		OutColor += AccumulateAndApplyWeight( SampleSceneColorRGBA(float2(Sample[2].x, Sample[2].y)), Weight[2].x * Weight[2].y, WeightsSum);
		OutColor += AccumulateAndApplyWeight( SampleSceneColorRGBA(float2(Sample[2].x, Sample[3].y)), Weight[2].x * Weight[3].y, WeightsSum);
		OutColor += AccumulateAndApplyWeight( SampleSceneColorRGBA(float2(Sample[2].x, Sample[4].y)), Weight[2].x * Weight[4].y, WeightsSum);

		OutColor += AccumulateAndApplyWeight( SampleSceneColorRGBA(float2(Sample[3].x, Sample[1].y)), Weight[3].x * Weight[1].y, WeightsSum);
		OutColor += AccumulateAndApplyWeight( SampleSceneColorRGBA(float2(Sample[3].x, Sample[2].y)), Weight[3].x * Weight[2].y, WeightsSum);
		OutColor += AccumulateAndApplyWeight( SampleSceneColorRGBA(float2(Sample[3].x, Sample[3].y)), Weight[3].x * Weight[3].y, WeightsSum);

		OutColor += AccumulateAndApplyWeight(SampleSceneColorRGBA(float2(Sample[4].x, Sample[2].y)), Weight[4].x * Weight[2].y, WeightsSum);

		OutColor /= WeightsSum;
	}

#elif METHOD == UPSCALE_METHOD_GAUSSIAN
	{
		// Gaussian filtered unsharp mask
		float2 UV = UVAndScreenPos.xy * Input_Extent;
		float2 tc = floor(UV) + 0.5;

		// estimate pixel value and derivatives
		OutColor = 0;
		float4 Laplacian = 0;
		float WeightsSum = 0;
		UNROLL for (int i = -3; i <= 2; ++i)
		{
			UNROLL for (int j = -3; j <= 2; ++j)
			{
				float2 TexelOffset = float2(i, j) + 0.5;

				// skip corners: eliminated entirely by UNROLL
				if (dot(TexelOffset, TexelOffset) > 9) continue;

				float2 Texel = tc + TexelOffset;
				float2 Offset = UV - Texel;
				float OffsetSq = 2 * dot(Offset, Offset);	// texel loop is optimized for variance = 0.5
				float Weight = exp(-0.5 * OffsetSq);
				float4 Sample = AccumulateAndApplyWeight(SampleSceneColorRGBA(Texel * Input_ExtentInverse), Weight, WeightsSum);

				OutColor += Sample;
				Laplacian += Sample * (OffsetSq - 2);
			}
		}

		const float InvWeightsSum = 1.0f / WeightsSum;
		OutColor  *= InvWeightsSum;
		Laplacian *= InvWeightsSum;

		float UnsharpScale = UpscaleSoftness * (1 - Input_Extent.x * Input_Extent.y * Output_ViewportSizeInverse.x * Output_ViewportSizeInverse.y);
		OutColor -= UnsharpScale * Laplacian;
	}

#elif METHOD == UPSCALE_METHOD_SMOOTHSTEP
	// smooth step.
	{
		OutColor = SampleSceneColorRGBA(GetSmoothstepUV(UVAndScreenPos.xy, Input_Extent, Input_ExtentInverse));
	}
#endif

#if defined(DIM_ALPHA_CHANNEL) && DIM_ALPHA_CHANNEL == 1
	FLATTEN
	if (OutColor.a > 0.999)
	{
		OutColor.a = 1;
	}
	else if (OutColor.a < 0.001)
	{
		OutColor.a = 0.0;
	}

	OutColor.a = select(bInvertAlpha > 0, 1.0f - OutColor.a, OutColor.a);
	
#else
	OutColor.a = 0; // Skip all computations related to alpha
#endif
}