// Copyright 1998-2017 Epic Games, Inc. All Rights Reserved.

/*=============================================================================
	PostProcessTemporalAA.usf: Temporal AA
=============================================================================*/

#include "Common.usf"
#include "PostProcessCommon.usf"
#include "Random.usf"
#include "EyeAdaptationCommon.usf"

#ifndef ENABLE_TEMPORAL_AA
#define ENABLE_TEMPORAL_AA	1
#endif

float SampleWeights[9];
float PlusWeights[5];
float VelocityScaling;
float CurrentFrameWeight;

#if COMPUTESHADER
// x: bForceResponsiveFrame, y: Unused, zw: ThreadToUVScale
float4 TemporalAAComputeParams;
RWTexture2D<float4> OutComputeTex;
#endif

float3 RGBToYCoCg( float3 RGB )
{
	float Y  = dot( RGB, float3(  1, 2,  1 ) );
	float Co = dot( RGB, float3(  2, 0, -2 ) );
	float Cg = dot( RGB, float3( -1, 2, -1 ) );
	
	float3 YCoCg = float3( Y, Co, Cg );
	return YCoCg;
}

float3 YCoCgToRGB( float3 YCoCg )
{
	float Y  = YCoCg.x * 0.25;
	float Co = YCoCg.y * 0.25;
	float Cg = YCoCg.z * 0.25;

	float R = Y + Co - Cg;
	float G = Y + Cg;
	float B = Y - Co - Cg;

	float3 RGB = float3( R, G, B );
	return RGB;
}

// Faster but less accurate luma computation. 
// Luma includes a scaling by 4.
float Luma4(float3 Color)
{
	return (Color.g * 2.0) + (Color.r + Color.b);
}

// Optimized HDR weighting function.
float HdrWeight4(float3 Color, float Exposure) 
{
	return rcp(Luma4(Color) * Exposure + 4.0);
}

float HdrWeightY(float Color, float Exposure) 
{
	return rcp(Color * Exposure + 1.0);
}

float HdrWeightG(float3 Color, float Exposure) 
{
	return rcp(Color.g * Exposure + 1.0);
}

float HdrWeightG_(float Color, float Exposure) 
{
	return rcp(Color * Exposure + 1.0);
}


// Optimized HDR weighting function.
float HdrWeight4_(float Color, float Exposure) 
{
	return rcp(Color * Exposure + 4.0);
}

// Optimized HDR weighting inverse.
float HdrWeightInv4(float3 Color, float Exposure) 
{
	return 4.0 * rcp(Luma4(Color) * (-Exposure) + 1.0);
}

float HdrWeightInvG(float3 Color, float Exposure) 
{
	return rcp(Color.g * (-Exposure) + 1.0);
}

float HdrWeightInvY(float Color, float Exposure) 
{
	return rcp(Color * (-Exposure) + 1.0);
}

float HdrWeightInv4_(float Color, float Exposure) 
{
	return 4.0 * rcp(Color * (-Exposure) + 1.0);
}

float HdrWeightInvG_(float Color, float Exposure) 
{
	return rcp(Color * (-Exposure) + 1.0);
}


// This returns exposure normalized linear luma from a PerceptualLuma4().
float LinearLuma4(float Channel, float Exposure) 
{
	return Channel * HdrWeightInv4_(Channel, Exposure);
}

// This returns exposure normalized linear luma from a PerceptualLuma4().
float LinearLumaG(float Channel, float Exposure) 
{
	return Channel * HdrWeightInvG_(Channel, Exposure);
}


float PerceptualLuma4(float3 Color, float Exposure) 
{
	float L = Luma4(Color);
	return L * HdrWeight4_(L, Exposure);
}

float PerceptualLumaG(float3 Color, float Exposure) 
{
	return Color.g * HdrWeightG_(Color.g, Exposure);
}



float Luma(float3 Color) 
{
	#if 1
		// This seems to work better (less same luma ghost trails).
		// CCIR 601 function for luma.
		return dot(Color, float3(0.299, 0.587, 0.114));
	#else
		// Rec 709 function for luma.
		return dot(Color, float3(0.2126, 0.7152, 0.0722));
	#endif
}

float HighlightCompression(float Channel) 
{
	return Channel * rcp(1.0 + Channel);
}
	
float HighlightDecompression(float Channel) 
{
	return Channel * rcp(1.0 - Channel);
}

float PerceptualLuma(float3 Color, float Exposure) 
{
	return sqrt(HighlightCompression(Luma(Color) * Exposure));
}

float LinearLuma(float Channel) 
{
	// This returns exposure normalized linear luma from a PerceptualLuma().
	return HighlightDecompression(Channel * Channel);
}

// Intersect ray with AABB, knowing there is an intersection.
//   Dir = Ray direction.
//   Org = Start of the ray.
//   Box = Box is at {0,0,0} with this size.
// Returns distance on line segment.
float IntersectAABB(float3 Dir, float3 Org, float3 Box)
{
	#if PS4_PROFILE
	    // This causes flicker, it should only be used on PS4 until proper fix is in.
		if(min(min(abs(Dir.x), abs(Dir.y)), abs(Dir.z)) < (1.0/65536.0)) return 1.0;
	#endif
	float3 RcpDir = rcp(Dir);
	float3 TNeg = (  Box  - Org) * RcpDir;
	float3 TPos = ((-Box) - Org) * RcpDir;
	return max(max(min(TNeg.x, TPos.x), min(TNeg.y, TPos.y)), min(TNeg.z, TPos.z));
}

float HistoryClamp(float3 History, float3 Filtered, float3 NeighborMin, float3 NeighborMax)
{
	float3 Min = min(Filtered, min(NeighborMin, NeighborMax));
	float3 Max = max(Filtered, max(NeighborMin, NeighborMax));	
	float3 Avg2 = Max + Min;
	float3 Dir = Filtered - History;
	float3 Org = History - Avg2 * 0.5;
	float3 Scale = Max - Avg2 * 0.5;
	return saturate(IntersectAABB(Dir, Org, Scale));	
}

float HdrWeight(float3 Color, float Exposure) 
{
	return rcp(max(Luma(Color) * Exposure, 1.0));
}

float4 HdrLerp(float4 ColorA, float4 ColorB, float Blend, float Exposure) 
{
	float BlendA = (1.0 - Blend) * HdrWeight(ColorA.rgb, Exposure);
	float BlendB =        Blend  * HdrWeight(ColorB.rgb, Exposure);
	float RcpBlend = rcp(BlendA + BlendB);
	BlendA *= RcpBlend;
	BlendB *= RcpBlend;
	return ColorA * BlendA + ColorB * BlendB;
}

void WeightTrackedAlphaClamping(
		inout float NeighborAlphaMin, inout float NeighborMinWeight,
		inout float NeighborAlphaMax, inout float NeighborMaxWeight,
		in float NeighborAlpha, in float NeighborWeight)
{
	if (NeighborAlpha < NeighborAlphaMin)
	{
		NeighborAlphaMin = NeighborAlpha;
		NeighborMinWeight = NeighborWeight;
	}
	else if (NeighborAlpha > NeighborAlphaMax)
	{
		NeighborAlphaMax = NeighborAlpha;
		NeighborMaxWeight = NeighborWeight;
	}
}

#define MINIMAL_REPROJECTION_DEBUG 0

// Useful for verifying that the camera motion reprojection is correct, and the history texture contains the right values
float3 MinimalReprojectionForDebugging(float2 UV, float4 SvPosition)
{
	float3 PosN; 
	PosN.xy = SvPositionToScreenPosition(SvPosition).xy;
	PosN.z = SceneDepthTexture.SampleLevel(SceneDepthTextureSampler, UV, 0).r;

	float4 ThisClip = float4( PosN.xy, PosN.z, 1 );
	float4 PrevClip = mul( ThisClip, View.ClipToPrevClip );
	float2 PrevScreen = PrevClip.xy / PrevClip.w;
	float2 BackN = PosN.xy - PrevScreen;

	BackN = SvPositionToScreenPosition(SvPosition).xy - BackN;
	bool OffScreen = max(abs(BackN.x), abs(BackN.y)) >= 1.0;
	BackN.x = clamp(BackN.x, -1.0 + ViewportSize.z, 1.0 - ViewportSize.z);
	BackN.y = clamp(BackN.y, -1.0 + ViewportSize.w, 1.0 - ViewportSize.w);
	BackN = (BackN * ScreenPosToPixel.xy + ScreenPosToPixel.zw + 0.5) * PostprocessInput0Size.zw;

	float3 OldColor = PostprocessInput1.SampleLevel(PostprocessInput1Sampler, BackN.xy, 0).rgb;
	float3 CurrentFrameSceneColor = PostprocessInput0.SampleLevel(PostprocessInput0Sampler, UV, 0).rgb;

	float FixedLerp = .1f;
	float3 NewColor = lerp(OldColor, CurrentFrameSceneColor, FixedLerp);
	return NewColor;
}

void SSRTemporalAAPS( 
	noperspective 	float2 UV 						: TEXCOORD0, 
	noperspective	float3 InExposureScaleVignette 	: TEXCOORD1, 
	noperspective	float4 Unused2 					: TEXCOORD2,
	noperspective	float4 Unused3 					: TEXCOORD3,
	noperspective	float2 Unused4 					: TEXCOORD4,
					float4 SvPosition 				: SV_POSITION,		// after all interpolators
	out 			float4 OutColor 				: SV_Target0 
	)
{
#if ENABLE_TEMPORAL_AA
	float InExposureScale = InExposureScaleVignette.x;
	#define AA_ALPHA 1
	#define AA_CROSS 0
	#define AA_DYNAMIC 1
	#define AA_LERP 8
	#define AA_NAN 1
	#include "PostProcessTemporalCommon.usf"
#else 
	// On broken platforms then at least draw something without AA.
	OutColor = PostprocessInput0.SampleLevel(PostprocessInput0Sampler, UV, 0);
#endif
}

void DOFTemporalAAPS( 
	in noperspective float2 UV : TEXCOORD0, 
	in noperspective float3 InExposureScaleVignette : TEXCOORD1, 
	in noperspective float4 Unused1 : TEXCOORD2,
	in noperspective float4 Unused2 : TEXCOORD3,
	in noperspective float2 Unused3 : TEXCOORD4,
	float4 SvPosition : SV_POSITION,		// after all interpolators
	out float4 OutColor : SV_Target0 )
{
#if ENABLE_TEMPORAL_AA
	float InExposureScale = InExposureScaleVignette.x;
	#define AA_FILTERED 0
	#define AA_ALPHA 1
	#define AA_CROSS 4
	#define AA_DYNAMIC 1
	#define AA_NAN 1
	#define AA_BORDER 1
	#define AA_FORCE_ALPHA_CLAMP 1
	#define AA_YCOCG 1
	#include "PostProcessTemporalCommon.usf"
#else 
	OutColor = PostprocessInput0.SampleLevel(PostprocessInput0Sampler, UV, 0);
#endif
}

#if COMPUTESHADER
[numthreads(THREADGROUP_SIZEX, THREADGROUP_SIZEY, 1)]
void DOFTemporalAACS(uint2 DispatchThreadId : SV_DispatchThreadID)
{
	float2 UV = ((float2)DispatchThreadId + View.ViewRectMin.xy + 0.5f) * TemporalAAComputeParams.zw;
	float4 OutColor = 0;

	if (IsComputeUVOutOfBounds(UV))
	{
		return;
	}

#if ENABLE_TEMPORAL_AA
	float InExposureScale = EyeAdaptationLookup();
	#define AA_FILTERED 0
	#define AA_ALPHA 1
	#define AA_CROSS 4
	#define AA_DYNAMIC 1
	#define AA_NAN 1
	#define AA_BORDER 1
	#define AA_FORCE_ALPHA_CLAMP 1
	#define AA_YCOCG 1
	#include "PostProcessTemporalCommon.usf"
#else
	OutColor = PostprocessInput0.SampleLevel(PostprocessInput0Sampler, UV, 0);
#endif

	uint2 PixelPos = DispatchThreadId + (uint2)View.ViewRectMin.xy;
	OutComputeTex[PixelPos] = OutColor;
}
#endif

void LightShaftTemporalAAPS( 
	in noperspective float2 UV : TEXCOORD0, 
	in noperspective float3 InExposureScaleVignette : TEXCOORD1, 
	in noperspective float4 Unused1 : TEXCOORD2,
	in noperspective float4 Unused2 : TEXCOORD3,
	in noperspective float2 Unused3 : TEXCOORD4,
	float4 SvPosition : SV_POSITION,		// after all interpolators
	out float4 OutColor : SV_Target0 )
{
#if ENABLE_TEMPORAL_AA
	float InExposureScale = InExposureScaleVignette.x;
	#define AA_ALPHA 1
	#define AA_CROSS 0
	#define AA_DYNAMIC 0
	#define AA_LERP 64
	#define AA_NAN 1
	#include "PostProcessTemporalCommon.usf"
#else 
	OutColor = PostprocessInput0.SampleLevel(PostprocessInput0Sampler, UV, 0);
#endif
}

// 0.0f or 1.0f (if the RT format is not PF_FloatRGBA)
float DitherScale;

float4 MainTemporalAACommon(in float2 UV, in float InExposureScale)
{
	float4 OutColor = 0;

	#define AA_FILTERED 1
	#define AA_ALPHA (POST_PROCESS_ALPHA)
	#define AA_BORDER 1
	#define AA_CROSS 2
	#define AA_AABB 1
	#define AA_YCOCG 1
	#define AA_BICUBIC 1
	#define AA_ACCURATE_ALPHA (POST_PROCESS_ALPHA)
	#define AA_DYNAMIC_ANTIGHOST 1
#if COMPUTESHADER
	#define AA_SINGLE_PASS_RESPONSIVE 1
#endif
	#include "PostProcessTemporalCommon.usf"

	return OutColor;
}

void MainTemporalAAPS(
	in noperspective float2 UV : TEXCOORD0,
	in noperspective float3 InExposureScaleVignette : TEXCOORD1,
	// Note: these are needed to match up with the vertex shader outputs which is FPostProcessTonemapVS, otherwise a d3d debug error is issued
	in noperspective float4 Unused1 : TEXCOORD2,
	in noperspective float4 Unused2 : TEXCOORD3,
	in noperspective float2 Unused3 : TEXCOORD4,
	in float4 SvPosition : SV_Position,		// after all interpolators
	out float4 OutColor : SV_Target0 )
{
#if ENABLE_TEMPORAL_AA
	#if MINIMAL_REPROJECTION_DEBUG
		OutColor = float4(MinimalReprojectionForDebugging(UV, SvPosition), 0);
	#else
		float InExposureScale = InExposureScaleVignette.x;
		OutColor = MainTemporalAACommon(UV, InExposureScale);
	#endif
#else
	OutColor = PostprocessInput0.SampleLevel(PostprocessInput0Sampler, UV, 0);
#endif
}

#if COMPUTESHADER
[numthreads(THREADGROUP_SIZEX, THREADGROUP_SIZEY, 1)]
void MainTemporalAACS(uint2 DispatchThreadId : SV_DispatchThreadID)
{
	float2 UV = ((float2)DispatchThreadId + View.ViewRectMin.xy + 0.5f) * TemporalAAComputeParams.zw;
	float4 SvPosition = float4(UV*2-1,0,1);
	float4 OutColor = 0;

	if (IsComputeUVOutOfBounds(UV))
	{
		return;
	}

#if ENABLE_TEMPORAL_AA
	#if MINIMAL_REPROJECTION_DEBUG
		OutColor = float4(MinimalReprojectionForDebugging(UV, SvPosition), 0);
	#else
		float InExposureScale = EyeAdaptationLookup();
		OutColor = MainTemporalAACommon(UV, InExposureScale);
	#endif
#else
	OutColor = PostprocessInput0.SampleLevel(PostprocessInput0Sampler, UV, 0);
#endif

	uint2 PixelPos = DispatchThreadId + (uint2)View.ViewRectMin.xy;
	OutComputeTex[PixelPos] = OutColor;
}
#endif

float4 MainFastTemporalAACommon(in float2 UV, in float InExposureScale)
{
	float4 OutColor = 0;

	#define AA_FILTERED 1
	#define AA_ALPHA (POST_PROCESS_ALPHA)
	#define AA_BORDER 1
	#define AA_CROSS 2
	#define AA_YCOCG 1
	#define AA_BICUBIC 1
	#define AA_ACCURATE_ALPHA (POST_PROCESS_ALPHA)
#if COMPUTESHADER
	#define AA_SINGLE_PASS_RESPONSIVE 1
#endif
	#include "PostProcessTemporalCommon.usf"

	return OutColor;
}

void MainFastTemporalAAPS(
	in noperspective float2 UV: TEXCOORD0,
	in noperspective float3 InExposureScaleVignette : TEXCOORD1,
	in noperspective float4 Unused1 : TEXCOORD2,
	in noperspective float4 Unused2 : TEXCOORD3,
	in noperspective float2 Unused3 : TEXCOORD4,
	in float4 SvPosition : SV_Position,		// after all interpolators
	out float4 OutColor : SV_Target0 )
{
#if ENABLE_TEMPORAL_AA
	float InExposureScale = InExposureScaleVignette.x;
	OutColor = MainFastTemporalAACommon(UV, InExposureScale);
#else
	OutColor = PostprocessInput0.SampleLevel(PostprocessInput0Sampler, UV, 0);
#endif
} 

#if COMPUTESHADER
[numthreads(THREADGROUP_SIZEX, THREADGROUP_SIZEY, 1)]
void MainFastTemporalAACS(uint2 DispatchThreadId : SV_DispatchThreadID)
{
	float2 UV = ((float2)DispatchThreadId + View.ViewRectMin.xy + 0.5f) * TemporalAAComputeParams.zw;
	float4 OutColor = 0;

	if (IsComputeUVOutOfBounds(UV))
	{
		return;
	}

#if ENABLE_TEMPORAL_AA
	float InExposureScale = EyeAdaptationLookup();
	OutColor = MainFastTemporalAACommon(UV, InExposureScale);
#else
	OutColor = PostprocessInput0.SampleLevel(PostprocessInput0Sampler, UV, 0);
#endif

	uint2 PixelPos = DispatchThreadId + (uint2)View.ViewRectMin.xy;
	OutComputeTex[PixelPos] = OutColor;
}
#endif