#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, normal discontinuities 
   and ssao convergence of a sample (the higher the more reliably
   has the sample a correct ssao value)
*/
float DiscontinuityFilter(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 centerPos = ReconstructSamplePos(colorsTex, 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;

	const float convergenceThresh = 200.0f;

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

		aoSample = tex2Dlod(ssaoTex, sampleTexCoord);
		
		// check spatial discontinuity
		// note: using the depth from the color texture is not 100% correct as depth was
		// not scaled with the interpolated view vector depth yet ...
		samplePos = ReconstructSamplePos(colorsTex, sampleTexCoord.xy, bl, br, tl, tr);
		//samplePos = ReconstructSamplePos(ssaoTex, sampleTexCoord.xy, bl, br, tl, tr);

		len = min(SqrLen(centerPos - samplePos), 1e2f);
		spatialFactor = 1.0f / max(len, 1e-3f);

		convergenceFactor = aoSample.y + 1.0f;

		//sampleNorm = tex2Dlod(normalsTex, sampleTexCoord).xyz;
		//normalFactor = max(step(.2f, dot(sampleNorm, centerNormal)), 1e-2f);

		// combine the weights
		w = convergenceFactor * 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. We assume that 
	we are using half resolution ssao for this version of the combineSsao
*/
pixel CombineSsaoHalfRes(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 float ssaoFilterRadius,
						 uniform float4x4 modelViewProj,
						 uniform float3 bl,
						 uniform float3 br,
						 uniform float3 tl,
						 uniform float3 tr,
						 uniform float w,
						 uniform float h
						 )
{
	pixel OUT;

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

#if 0 // use half size

	// the following has to be done for half resolution ssao:
	// get the minimum convergence by exactly sampling the 4 surrounding
	// texels in the old texture, otherwise flickering because convergence
	// will be interpolated when upsampling and filter size does not match!

	float4 texelCenterConv;
	const float xoffs = .5f / w; const float yoffs = .5f / h;

	// get position exactly between old texel centers
	float2 center;
	center.x = (floor(IN.texCoord.x * w - .5f) + 1.0f) / w;
	center.y = (floor(IN.texCoord.y * h - .5f) + 1.0f) / h;

	texelCenterConv.x = tex2Dlod(ssaoTex, float4(center + float2( xoffs,  yoffs), 0, 0)).y;
	texelCenterConv.y = tex2Dlod(ssaoTex, float4(center + float2( xoffs, -yoffs), 0, 0)).y;
	texelCenterConv.z = tex2Dlod(ssaoTex, float4(center + float2(-xoffs, -yoffs), 0, 0)).y;
	texelCenterConv.w = tex2Dlod(ssaoTex, float4(center + float2(-xoffs,  yoffs), 0, 0)).y;

	const float m1 = min(texelCenterConv.x, texelCenterConv.y);
	const float m2 = min(texelCenterConv.z, texelCenterConv.w);

	const float minConvergence = min(m1, m2);
	const float convergence = minConvergence;

#else
	
	// just take unfiltered convergence in current pixel
	const float convergence = ao.y;

#endif
	
	// filter reaches size 1 pixel when sample size reaches threshold 
	// afterwards we do not use the filter anymore

	// filter up to a certain convergance value and leave out background (sky) by checking depth
	if ((convergence < SSAO_CONVERGENCE_THRESHOLD) && (col.w < 1e10f))
	{
		const float distanceScale = 1.0f;

		// descend to zero filter size after reaching thres pixels
		const float convergenceWeight = SSAO_CONVERGENCE_THRESHOLD / (ssaoFilterRadius - 1.0f);
		const float convergenceScale = convergenceWeight / (convergence + convergenceWeight);
		const float scale = ssaoFilterRadius * convergenceScale * distanceScale;

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

	// just apply ssao if we are not in the sky
	if (col.w < 1e10f)
		OUT.illum_col.xyz = col.xyz * max(2e-2f, 1.0f - ao.x);
		//OUT.illum_col.xyz = col.xyz * ao.x;
	else
		OUT.illum_col.xyz = col.xyz;


	//OUT.illum_col.xyz = float3(abs(ao.y * 1e2f), abs(ao.z * 1e2f), abs(ao.w * 1e2f));

	//if (convergence < (1.0f + NUM_SAMPLES * 10))
	//	OUT.illum_col.xyz = float3(1 - convergence / (NUM_SAMPLES * 10), convergence / (NUM_SAMPLES * 10), 0);
	
	//OUT.illum_col.xyz = float3(ao.z * 1e4f, ao.z * 1e4f, ao.z * 1e4f);

	//OUT.illum_col.xyz = float3(ao.x, ao.x, step(thres, convergence));
	//OUT.illum_col.xyz = float3(abs(center.x - IN.texCoord.x) * 16.0f, abs(center.y - IN.texCoord.y) * 12.0f, 0);
	//OUT.illum_col.xyz = float3(0, 1.0f - step(0.5f + NUM_SAMPLES, convergence), 1);

	OUT.illum_col.w = col.w;

	return OUT;
}