/*=============================================================================
	DOFAndBloomGatherPixelShader.usf: Pixel shader for gathering the combined depth of field and bloom samples for blurring.
	Copyright 1998-2010 Epic Games, Inc. All Rights Reserved.
=============================================================================*/

#include "Common.usf"

SAMPLER2D(DOFTexture);
SAMPLER2D(DOFTempTexture);
SAMPLER2D(DOFTempTexture2);

/** The amount bloomed colors are scaled by. */
half2 BloomScaleAndThreshold;

/** Computes a pixel's luminance for bloom */
half ComputeLuminanceForBloom( half3 InSceneColor )
{
	// Compute the luminance for this pixel
	half TotalLuminance;
	if( 1 )
	{
		// Compute luminance as the maximum of RGB.  This is a bit more intuitive for artists as they know
		// that any pixel with either of the RGB channels above 1.0 will begin to bloom.
		TotalLuminance = max( InSceneColor.r, max( InSceneColor.g, InSceneColor.b ) );
	}
	else
	{
		// RGB scale factor to calculated pixel luminance using a weight average
		half3 LuminanceFactor = half3( 0.3, 0.59, 0.11 );

		// Compute true luminance
		TotalLuminance = dot( LuminanceFactor, InSceneColor );
	}

	return TotalLuminance;
}



/** Computes bloomed amount for the specified scene color */
half ComputeBloomAmount( half3 InSceneColor, half InLuminance )
{
	// ME3START - Rob Krajcarski - October 2, 2011 - Preport of MotionBlur improvements (see UberPostProcessEffect.cpp for more details)
	// How bright a pixel must be (luminance) before it begins to bloom
	// NOTE: We could make this artist-configurable
	//half BloomThreshold = 1.0f;
	// ME3END

	// Size of the bloom "ramp".  This value specifies the amount of light beyond the bloom threshold required
	// before a pixel's bloom will be 100% of the original color.
	// NOTE: Any value above 0.8 looks pretty good here (and 1.0 is often fastest), but a value of 2.0 here
	//	     minimizes artifacts: the bloom ramp-up will closely match the linear ascent of additive color
	half BloomRampSize = 2.0f;

	// Figure out how much luminance is beyond the bloom threshold.  Note that this value could be negative but
	// we handle that in the next step.
	// ME3START - Rob Krajcarski - October 2, 2011 - Preport of MotionBlur improvements (see UberPostProcessEffect.cpp for more details)
	half BloomLuminance = InLuminance - BloomScaleAndThreshold.y;
	// ME3END

	// Note that we clamp the bloom amount between 0.0 and 1.0, but pixels beyond our bloom ramp will still 
	// bloom brighter because we'll use 100% of the original scene color as bloom
	half BloomAmount = saturate( BloomLuminance / BloomRampSize );
	
	return BloomAmount;
}



/** Computes bloomed amount for the specified scene color */
half ComputeBloomAmount( half3 InSceneColor )
{
	// Compute the luminance for this pixel
	half TotalLuminance = ComputeLuminanceForBloom( InSceneColor );

	return ComputeBloomAmount( InSceneColor, TotalLuminance );
}



/** Computes bloomed color for the specified scene color */
half3 ComputeBloomColor( half3 InSceneColor, half InLuminance )
{
	// Multiply with the scene color to get the final bloom amount
	return InSceneColor * ComputeBloomAmount( InSceneColor, InLuminance );
}


/** Computes bloomed color for the specified scene color */
half3 ComputeBloomColor( half3 InSceneColor )
{
	// Multiply with the scene color to get the final bloom amount
	return InSceneColor * ComputeBloomAmount( InSceneColor );
}


#include "DepthOfFieldCommon.usf"

/** The number of float4s the 2D sample offsets are packed into. */
#define NUM_CHUNKS ((NUM_SAMPLES + 1) / 2)

void CommonDOFSetup(in float2 CenterUV, out bool bFrontLayer, out float4 Front, out float4 Back )
{
	float2 Offset = InputTextureSize.zw;

	float2 UV[4];

	// no filtering (2x2 kernel) to get no leaking in Depth of Field
	UV[0] = CenterUV + Offset * float2(-0.5f, -0.5f);
	UV[1] = CenterUV + Offset * float2( 0.5f, -0.5f);
	UV[2] = CenterUV + Offset * float2(-0.5f,  0.5f);
	UV[3] = CenterUV + Offset * float2( 0.5f,  0.5f);

	float3 Color[4];
	float2 Layer[4];
	
	const float FrontGamma = 0.1f;
	float3 FrontColor = float3(0,0,0);

	UNROLL
	for(uint i = 0; i < 4; ++i)
	{
		Color[i] = CalcSceneColor(UV[i]);
		FrontColor += pow(Color[i], FrontGamma);

		//const float Depth = CalcSceneColorAndDepth(UV[i]).a;
		//const float Depth = 1.f / tex2D(DOFTexture, UV[i]).r;
		const float Depth = ConvertFromDeviceZ(saturate(tex2D(DOFTexture, UV[i]).r));

		Layer[i].x = ComputeNearCoC(Depth);
		Layer[i].y = ComputeFarCoC(Depth);
	}

	float2 LayerSum = Layer[0] + Layer[1] + Layer[2] + Layer[3];
	
	Front.rgb = pow(FrontColor * 0.25f, 1.f / FrontGamma);
	Front.a = (Layer[0].x + Layer[1].x + Layer[2].x + Layer[3].x) * .25f;

	float4 BackMask = float4(Layer[0].y, Layer[1].y, Layer[2].y, Layer[3].y);
	float SumBackMask = dot(BackMask, 1);
	if(SumBackMask > 0.001f)
	{
		Back.rgb = (
			Color[0] * BackMask.x +
			Color[1] * BackMask.y +
			Color[2] * BackMask.z +
			Color[3] * BackMask.w ) / SumBackMask;
	}
	else
	{
		Back.rgb = Color[0];
	}
	Back.a = min(Layer[0].y, min(Layer[1].y, min(Layer[2].y, Layer[3].y)));

	bFrontLayer = Front.a > Back.a;
}

/**
 * Entry point for the gather pass, downsample to half size and separate into near, far and a circle of confusion tilemap.
 */
void MainDOFOnly(
	in float4 OffsetUVs[NUM_CHUNKS] : TEXCOORD0,
	out float4 OutNearColor : SV_Target0,
	out float4 OutFarColor : SV_Target1
	)
{
	bool bFrontLayer;
	CommonDOFSetup( OffsetUVs[0].xy, bFrontLayer, OutNearColor, OutFarColor );
	
	float2 CoC = float2(OutNearColor.a, OutFarColor.a);
	CoC *= MinMaxBlurClamp.xy;
	CoC.y = max(CoC.x, CoC.y);
	//const float2 BlendOffset = float2(0.1f, 0.1f);
	//CoC = saturate(CoC - BlendOffset);
	CoC *= CoC;

	OutNearColor.rgb = min(float3(256 * 256, 256 * 256, 256 * 256), OutNearColor.rgb);
	OutFarColor.rgb = min(float3(256 * 256, 256 * 256, 256 * 256), OutFarColor.rgb);
	
	OutNearColor.a = CoC.x;
	OutFarColor.a = CoC.y;
}

float4 BlurDirections;

void MainDOFBlur1(
	in float4 OffsetUVs[NUM_CHUNKS] : TEXCOORD0,
	out float4 OutVertBlurred : SV_Target0,
	out float4 OutDiagBlurred : SV_Target1
	)
{
	const float2 uv = OffsetUVs[0].xy;
	const float invStep = 1.f / NUM_SAMPLES;
	const float4 stepUV = BlurDirections * DOFKernelParams.x * invStep;
	
	float4 baseColor = tex2D(DOFTempTexture, uv);
	float baseCoC = baseColor.a;

	float4 finalColor1 = float4(0,0,0,0);
	float4 finalColor2 = float4(0,0,0,0);
	
    float maxCoC1 = baseCoC;
    float maxCoC2 = baseCoC;

	const float4 baseUV = uv.xyxy + BlurDirections * 0.5f;

	UNROLL
	for (int i = 0; i < NUM_SAMPLES; ++i)
	{
		float4 sampleUV = baseUV + stepUV * i * baseCoC;
		float4 color1 = tex2D(DOFTempTexture, sampleUV.xy);
		float4 color2 = tex2D(DOFTempTexture, sampleUV.zw);
		
		maxCoC1 = max(maxCoC1, color1.a);
		maxCoC2 = max(maxCoC2, color2.a);

		float4 scale = float4(color1.a, color1.a, color2.a, color2.a);
		sampleUV = baseUV + stepUV * i * scale;
		color1 = tex2D(DOFTempTexture, sampleUV.xy);
		color2 = tex2D(DOFTempTexture, sampleUV.zw);
		
		color1.rgb *= color1.a;
		color2.rgb *= color2.a;

		finalColor1 += color1;
		finalColor2 += color2;
	}

	if (finalColor1.a >= 0.001f) finalColor1.rgb /= finalColor1.a;
	if (finalColor2.a >= 0.001f) finalColor2.rgb /= finalColor2.a;

	OutVertBlurred.rgb = finalColor1.rgb;
	OutDiagBlurred.rgb = finalColor1.rgb + finalColor2.rgb;

	OutVertBlurred.a = maxCoC1;
	OutDiagBlurred.a = max(maxCoC1, maxCoC2);
}

void MainDOFBlur2(
	in float4 OffsetUVs[NUM_CHUNKS] : TEXCOORD0,
	out float4 OutBlurred : SV_Target0
	)
{
	const float2 uv = OffsetUVs[0].xy;
	const float invStep = 1.f / NUM_SAMPLES;
	const float4 stepUV = BlurDirections * DOFKernelParams.x * invStep;
	
	float4 baseColor = tex2D(DOFTempTexture, uv);
	float baseCoC = baseColor.a;
	
	float4 finalColor1 = float4(0,0,0,0);
	float4 finalColor2 = float4(0,0,0,0);
		
    float maxCoC = baseCoC;
	
	const float4 baseUV = uv.xyxy + BlurDirections * 0.5f;
	
	UNROLL
	for (int i = 0; i < NUM_SAMPLES; ++i)
	{
		float4 sampleUV = baseUV + stepUV * i * baseCoC;
		float4 color1 = tex2D(DOFTempTexture, sampleUV.xy);
		float4 color2 = tex2D(DOFTempTexture2, sampleUV.zw);
		
		maxCoC = max(maxCoC, max(color1.a, color2.a));
		
		float4 scale = float4(color1.a, color1.a, color2.a, color2.a);
		sampleUV = baseUV + stepUV * i * scale;
		color1 = tex2D(DOFTempTexture, sampleUV.xy);
		color2 = tex2D(DOFTempTexture2, sampleUV.zw);
		
		color1.rgb *= color1.a;
		color2.rgb *= color2.a;

		finalColor1 += color1;
		finalColor2 += color2;
	}
	
	if (finalColor1.a >= 0.001f) finalColor1.rgb /= finalColor1.a;
	if (finalColor2.a >= 0.001f) finalColor2.rgb /= finalColor2.a;

	OutBlurred.rgb = (finalColor1.rgb + finalColor2.rgb) * .33f;
	OutBlurred.a = maxCoC;
}

void DOFBoxBlurMin(
	in float4 OffsetUVs[NUM_CHUNKS] : TEXCOORD0,
	out float4 OutBlurred : SV_Target0
	)
{
	const float2 uv = OffsetUVs[0].xy;

	OutBlurred = tex2D(DOFTempTexture, uv);
	
	const float invStep = 2.f / NUM_SAMPLES;
	const float2 stepUV = BlurDirections.xy * DOFKernelParams.x * 0.5f;

	float CoC = 0.f;

	UNROLL
	for (float x = -1.; x <= 1.; x += invStep)
	{
		CoC += min(tex2D(DOFTempTexture, uv + stepUV * x).a, OutBlurred.a);
	}
	CoC /= NUM_SAMPLES + 1;

	OutBlurred.a = CoC;
}

void DOFBoxBlurMax(
	in float4 OffsetUVs[NUM_CHUNKS] : TEXCOORD0,
	out float4 OutBlurred : SV_Target0
	)
{
	const float2 uv = OffsetUVs[0].xy;

	OutBlurred = tex2D(DOFTempTexture, uv);
	
	const float invStep = 2.f / NUM_SAMPLES;
	const float2 stepUV = BlurDirections.xy * DOFKernelParams.x * 0.5f;

	float CoC = 0.f;

	UNROLL
	for (float x = -1.; x <= 1.; x += invStep)
	{
		CoC += max(tex2D(DOFTempTexture, uv + stepUV * x).a, OutBlurred.a);
	}
	CoC /= NUM_SAMPLES + 1;

	OutBlurred.a = CoC;
}


/**
 * Entry point for the gather pass, which downsamples from scene color to the filter buffer. 
 * Bloom color is stored in OutColor.rgb.
 */
void MainBloomOnly(
	in float4 OffsetUVs[NUM_CHUNKS] : TEXCOORD0,
	// BIOSTART - September 15, 2010 - Wihlidal, Graham - PS3 Fragment and Vertex Shader Optimizations
	out half4 OutColor : SV_Target0
	// BIOEND
	)
{
	// ME3START - Rob Krajcarski - October 2, 2011 - Preport of MotionBlur improvements (see UberPostProcessEffect.cpp for more details)
	half3 AvgBloomColor = 0; 

	//Go through each chunk and take samples.  NUM_SAMPLES must be a factor of 2.
	for(int ChunkIndex = 0;ChunkIndex < NUM_SAMPLES / 2;ChunkIndex++)
	{
		// Sample scene color/depth (1) and accumulate average
		half3 SceneColor1 = CalcSceneColor(OffsetUVs[ChunkIndex].xy);
		
		// Sample scene color/depth (2) and accumulate average
		half3 SceneColor2 = CalcSceneColor(OffsetUVs[ChunkIndex].wz);

		// The bloom color is the scaled scene color if it has a component outside the displayable range [0,1].
		// Only bloom if (SceneColor > 1), instead of (0 > SceneColor > 1), in order to mimic XBOX behavior due to having unsigned SceneColor values
		// this comparison is done per scene color sample to reduce aliasing on high frequency bright patterns
		
		// Compute bloom amount (1) and accumulate average
		AvgBloomColor += ComputeBloomColor( SceneColor1 );

		// Compute bloom amount (2) and accumulate average
		AvgBloomColor += ComputeBloomColor( SceneColor2 );
	}

	//normalize and scale
	AvgBloomColor = AvgBloomColor * BloomScaleAndThreshold.x / NUM_SAMPLES;

	//scale output down to fit in the [0-1] range of the fixed point filter buffer, alpha is unused
	// BIOSTART - September 15, 2010 - Wihlidal, Graham - PS3 Fragment and Vertex Shader Optimizations
	OutColor = half4(AvgBloomColor, 0.0f) / MAX_SCENE_COLOR;

// ME3END
}

void MainDOFAndBloom(
 	in float4 OffsetUVs[NUM_CHUNKS] : TEXCOORD0,
 	out half4 OutColor : SV_Target0
 	)
{
	// RGB:scene color, A:depth
	half4 AvgSceneColorAndDepth = 0;
	half3 AvgBloomColor = 0; 

	//Go through each chunk and take samples.  NUM_SAMPLES must be a factor of 2.
	for(int ChunkIndex = 0;ChunkIndex < NUM_SAMPLES / 2;ChunkIndex++)
	{
		// Sample scene color/depth (1) and accumulate average
		half4 SceneColorAndDepth1 = CalcSceneColorAndDepth(OffsetUVs[ChunkIndex].xy);
		AvgSceneColorAndDepth += SceneColorAndDepth1;
		
		// Sample scene color/depth (2) and accumulate average
		half4 SceneColorAndDepth2 = CalcSceneColorAndDepth(OffsetUVs[ChunkIndex].wz);
		AvgSceneColorAndDepth += SceneColorAndDepth2;

		// The bloom color is the scaled scene color if it has a component outside the displayable range [0,1].
		// Only bloom if (SceneColor > 1), instead of (0 > SceneColor > 1), in order to mimic XBOX behavior due to having unsigned SceneColor values
		// this comparison is done per scene color sample to reduce aliasing on high frequency bright patterns

		// Compute bloom amount (1) and accumulate average
		AvgBloomColor += ComputeBloomColor( SceneColorAndDepth1 );

		// Compute bloom amount (2) and accumulate average
		AvgBloomColor += ComputeBloomColor( SceneColorAndDepth2 );
	}

	//normalize and scale
	AvgBloomColor = AvgBloomColor * BloomScaleAndThreshold.x / NUM_SAMPLES;
	AvgSceneColorAndDepth = AvgSceneColorAndDepth / NUM_SAMPLES;

	half UnfocusedPercent = CalcUnfocusedPercent(AvgSceneColorAndDepth.a);

	//scale output down to fit in the [0-1] range of the fixed point filter buffer
	OutColor = half4(AvgSceneColorAndDepth.rgb * UnfocusedPercent + AvgBloomColor, UnfocusedPercent) / MAX_SCENE_COLOR;
	
#if PS3
	// ME3START - Rob Krajcarski - November 21, 2011 - Fix for super HDR content on the PS3, which needs an explicit saturate()
	OutColor = isnan(OutColor) ? half4(0,0,0,0) : saturate(OutColor);
	//OutColor = isnan(OutColor) ? half4(0,0,0,0) : OutColor;
	// ME3END
#endif
}

// ME3END