#include "../shaderenv.h"
#include "common.h"


/*************************************************/
/*     Filter for combining ssao with image      */
/*************************************************/


struct fragment
{
	float2 texCoord: TEXCOORD0; 
	float3 view: TEXCOORD1;
};


struct pixel
{
	float4 illum_col: COLOR0;
};


/** Filter taking into account depth and normal differences, 
    and convergence of a sample  
*/
float DiscontinuityFilter(float2 texCoord,
						  float4 ao,
						  float4 color,
						  uniform sampler2D ssaoTex,
						  uniform sampler2D normalsTex,
						  uniform sampler2D colorsTex,
						  uniform float2 filterOffs[NUM_SSAO_FILTER_SAMPLES],
						  uniform float filterWeights[NUM_SSAO_FILTER_SAMPLES],
						  float scale
						  )
{
	float average = .0f;
	float total_w = .0f;

	const float eyeSpaceDepth = color.w;

	const float3 centerNormal = tex2Dlod(normalsTex, float4(texCoord, 0, 0)).xyz;

	float4 aoSample;
	float3 sampleNorm;
	float3 samplePos;
	float w;
	float4 sampleTexCoord;
	float depthFactor;
	float normalFactor;
	float convergenceFactor;
	float sampleDepth;

	for (int i = 0; i < NUM_SSAO_FILTER_SAMPLES; ++ i)
	{
		sampleTexCoord = float4(texCoord + filterOffs[i] * scale, 0, 0);

		aoSample = tex2Dlod(ssaoTex, sampleTexCoord);
		sampleNorm = tex2Dlod(normalsTex, sampleTexCoord).xyz;

		// check depth discontinuity
		sampleDepth = tex2Dlod(colorsTex, sampleTexCoord).w;
		//sampleDepth = aoSample.w;
		
		//depthFactor = 1.0f / max(abs(eyeSpaceDepth - sampleDepth), 1e-2f);
		depthFactor = 1.0f - step(1e-2f, abs(1.0f - eyeSpaceDepth / sampleDepth));
		//normalFactor = max(step(0.6f, dot(sampleNorm, centerNormal)), 1e-3f);
		normalFactor = max(dot(sampleNorm, centerNormal), 1e-3f);
		//convergenceFactor = min(100.0f, aoSample.y);
		convergenceFactor = aoSample.y;

		// combine the weights
		w = filterWeights[i] * convergenceFactor * depthFactor * normalFactor;
		//w = normalFactor * convergenceFactor;

		average += aoSample.x * w;
		total_w += w;
	}

	average /= max(total_w, 1e-6f);

	return saturate(average);
}



/** Filter taking into account depth and normal differences, 
   and convergence of a sample  
*/
float DiscontinuityFilter2(float2 texCoord,
						   float4 ao,
						   float4 color,
						   uniform sampler2D ssaoTex,
						   uniform sampler2D normalsTex,
						   uniform sampler2D colorsTex,
						   uniform float2 filterOffs[NUM_SSAO_FILTER_SAMPLES],
						   float scale,
						   float3 bl,
						   float3 br,
						   float3 tl,
						   float3 tr)
{
	float average = .0f;
	float total_w = .0f;

	const float3 centerPos = ReconstructSamplePos(ssaoTex, texCoord, bl, br, tl, tr);
	const float3 centerNormal = tex2Dlod(normalsTex, float4(texCoord, 0, 0)).xyz;

	float4 aoSample;
	float3 sampleNorm;
	float3 samplePos;
	float w;
	float4 sampleTexCoord;
	float spatialFactor;
	float normalFactor;
	float convergenceFactor;
	float len;

	for (int i = 0; i < NUM_SSAO_FILTER_SAMPLES; ++ i)
	{
		sampleTexCoord = float4(texCoord + filterOffs[i] * scale, 0, 0);

		aoSample = tex2Dlod(ssaoTex, sampleTexCoord);
		sampleNorm = tex2Dlod(normalsTex, sampleTexCoord).xyz;

		// check spatial discontinuity
		samplePos = ReconstructSamplePos(ssaoTex, sampleTexCoord.xy, bl, br, tl, tr);
		len = min(SqrLen(centerPos - samplePos), 1e2f);

		spatialFactor = 1.0f / max(len, 1e-3f);

		//normalFactor = max(step(0.6f, dot(sampleNorm, centerNormal)), 1e-3f);
		normalFactor = max(dot(sampleNorm, samplePos), 1e-3f);
		//convergenceFactor = min(100.0f, aoSample.y);
		convergenceFactor = aoSample.y;

		// combine the weights
		w = convergenceFactor * spatialFactor * normalFactor;

		average += aoSample.x * w;
		total_w += w;
	}

	average /= max(total_w, 1e-6f);

	return saturate(average);
}


/** Function combining image and indirect illumination buffer using a 
	depth and normal aware discontinuity filter.
*/
pixel combine(fragment IN, 
			  uniform sampler2D colorsTex,
			  uniform sampler2D ssaoTex,
			  uniform sampler2D normalsTex,
			  uniform float2 filterOffs[NUM_SSAO_FILTER_SAMPLES],
			  uniform float filterWeights[NUM_SSAO_FILTER_SAMPLES],
			  uniform float4x4 modelViewProj,
			  uniform float3 bl,
			  uniform float3 br,
			  uniform float3 tl,
			  uniform float3 tr
			  )
{
	pixel OUT;

	float4 col = tex2Dlod(colorsTex, float4(IN.texCoord, 0, 0));
	float4 ao = tex2Dlod(ssaoTex, float4(IN.texCoord, 0, 0));

	// reconstruct position from the eye space depth
	/*const float4 worldPos = float4(-IN.view * eyeSpaceDepth, 1.0f);
	// compute w factor from projection in order to control filter size
	const float4 projPos = mul(modelViewProj, worldPos);
	const float distanceScale = 1.0f / projPos.w;
*/

	// filter up to a certain convergance value and leave out background (sky) by checking depth
	if ((ao.y < 100.0f) && (col.w < 1e10f))
	//if (col.w < 1e10f)
	{
		//const float eyeSpaceDepth = col.w; const float distanceScaleWeight = 2.0f;
		//const float distanceScale = distanceScaleWeight / (eyeSpaceDepth + distanceScaleWeight);
		const float distanceScale = 1.0f;

		const float convergence = ao.y;
		const float adaptWeight = 25.0f;
		const float convergenceScale = SSAO_CONVERGENCE_WEIGHT / (convergence + SSAO_CONVERGENCE_WEIGHT);

		const float scale = NUM_SSAO_FILTER_WIDTH * convergenceScale * distanceScale;

		// the filtered ssao value
		ao.x = DiscontinuityFilter2(IN.texCoord, ao, col, ssaoTex, normalsTex, colorsTex, filterOffs, scale, bl, br, tl, tr);
	}

	OUT.illum_col.xyz = col.xyz * ao.x;
	//OUT.illum_col.xyz = float3(ao.x, ao.x, ao.x);
	//OUT.illum_col.xyz = float3(0, clamp(1.0f - ao.y * 1e-2f, 0, 1), 1);
	OUT.illum_col.w = col.w;

	return OUT;
}