////////////////////
// SSAO + color bleeding shader
// based on shader of Alexander Kusternig

#include "../shaderenv.h"


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


struct pixel2
{
	float4 ssao_col: COLOR0;
	float4 illum_col: COLOR1;
};


struct pixel
{
	float4 illum_col: COLOR0;
};


float2 myreflect(float2 pt, float2 n)
{
	// distance to plane
	float d = dot(n, pt);
	// reflect around plane
	float2 rpt = pt - d * 2.0f * n;
	return rpt;
}

struct GiStruct
{
	float3 illum;
	float2 ao;
};



inline float3 Interpol(float2 w, float3 bl, float3 br, float3 tl, float3 tr)
{
	float3 x1 = lerp(bl, tl, w.y);
	float3 x2 = lerp(br, tr, w.y); 
	float3 v = lerp(x1, x2, w.x); 

	return v;
}


/** Computes  diffuse reflections + ambient occlusion
*/
GiStruct globIllum(fragment IN,
				   uniform sampler2D colors,
				   uniform sampler2D noiseTexture,
				   uniform float2 samples[NUM_SAMPLES],
				   uniform float3 currentNormal,
				   uniform float3 centerPosition,
				   float scaleFactor,
				   uniform float3 bl,
				   uniform float3 br,
				   uniform float3 tl,
				   uniform float3 tr
				   , uniform float3 viewDir
				   )
{
	GiStruct gi;

	// Check in a circular area around the current position.
	// Shoot vectors to the positions there, and check the angle to these positions.
	// Summing up these angles gives an estimation of the occlusion at the current position.

	// ao is in stored in the w component
	float3 total_color = float3(0, 0, 0);
	float total_ao = 0.0f;
	float numSamples = 0.0f;

	////////////
	//-- the main sampling loop

	for (int i = 0; i < NUM_SAMPLES; i ++) 
	{
		const float2 offset = samples[i];

#if 1
		////////////////////
		// add random noiseTex: reflect around random normal vector (warning: slow!)
		float2 mynoise = tex2D(noiseTexture, IN.texCoord.xy).xy;
		float2 offsetTransformed = myreflect(offset, mynoise);
#else
		float2 offsetTransformed = offset;
#endif
		// weight with projected coordinate to reach similar kernel size for near and far
		float2 texcoord = IN.texCoord.xy + offsetTransformed * scaleFactor;

		//if ((texcoord.x <= 1.0f) && (texcoord.x >= 0.0f) && (texcoord.y <= 1.0f) && (texcoord.y >= 0.0f)) ++ numSamples;


		//////////
		//-- reconstruct world space position from sample

		const float4 sample = tex2Dlod(colors, float4(texcoord, 0, 0));
		//const float4 sample = tex2D(colors, texcoord);

		const float eyeSpaceDepth = sample.w;
		float3 rotView = Interpol(texcoord, bl, br, tl, tr);
		
		const float3 samplePos = - rotView * eyeSpaceDepth;
		const float3 sampleCol = sample.xyz;

		// distance between current position and sample position controls AO intensity.

		float3 dirSample = samplePos - centerPosition.xyz;
		const float magSample = length(dirSample);
		// normalize
		dirSample /= magSample;

		// use angle between current normal and direction to sample controls AO intensity.
		float cosAngle = max(dot(dirSample, currentNormal), 0);

		const float denom = (DISTANCE_SCALE + magSample * magSample);

		float2 intensity = float2(SAMPLE_INTENSITY, ILLUM_INTENSITY);
		intensity /= denom;

		// if normal perpenticular to view dir, only the samples approx count half
#if 1
		const float viewCorrect = 1.0f + VIEW_CORRECTION_SCALE * dot(viewDir, currentNormal);
	
		total_ao += cosAngle * intensity.x * viewCorrect;
		total_color += cosAngle * intensity.y * sampleCol * viewCorrect;
#else
		total_ao += cos_angle * intensity.x;
		total_color += cos_angle * intensity.y * sampleCol;
#endif
	}

	gi.illum = total_color;
	gi.ao = float2(max(0.0f, 1.0f - total_ao), numSamples);

	return gi;
}



pixel2 main(fragment IN, 
		   uniform sampler2D colors,
		   uniform sampler2D normals,
		   uniform sampler2D noiseTex,
		   uniform float2 samples[NUM_SAMPLES],
		   uniform sampler2D oldSsaoTex,
		   uniform sampler2D oldIllumTex,
		   const uniform float4x4 oldModelViewProj,
		   const uniform float4x4 modelViewProj,
		   uniform float temporalCoherence,
		   uniform float3 bl,
		   uniform float3 br,
		   uniform float3 tl,
		   uniform float3 tr
		   )
{
	pixel2 OUT;

	float4 norm = tex2Dlod(normals, float4(IN.texCoord, 0 ,0));
	float3 normal = normalize(norm.xyz);	

	
	/////////////
	//-- reconstruct position from the eye space depth

	float3 viewDir = IN.view;
	const float eyeDepth = tex2Dlod(colors, float4(IN.texCoord, 0, 0)).w;
	
	const float4 eyeSpacePos = float4(-viewDir * eyeDepth, 1.0f);

	// calculcate the current projected depth for next frame
	float4 currentPos = mul(modelViewProj, eyeSpacePos);

	const float w = SAMPLE_RADIUS / currentPos.w;
	currentPos /= currentPos.w;
	
	const float currentDepth = currentPos.z * PRECISION_SCALE;

	///////////
	//-- compute color bleeding + ao

	GiStruct gi = globIllum(IN, colors, noiseTex, samples, normal, eyeSpacePos, w, bl, br, tl, tr, normalize(IN.view));
	

	/////////////////
	//-- compute temporally smoothing

	// reprojection new frame into old one 
	// calculate projected depth
	float4 projPos = mul(oldModelViewProj, eyeSpacePos);
	projPos /= projPos.w;

	// the current depth projected into the old frame
	const float projDepth = projPos.z * PRECISION_SCALE;

	// fit from unit cube into 0 .. 1
	float2 tex = projPos.xy * 0.5f + 0.5f;

	// retrieve the sample from the last frame
	float3 oldSsao = tex2D(oldSsaoTex, tex).xyz;
	float3 oldIllum = tex2D(oldIllumTex, tex).xyz;

	const float oldDepth = oldSsao.z;
	//const float depthDif = 1.0f - projDepth / oldDepth;
	const float depthDif = projDepth - oldDepth;

	// the weights that indicate the state of convergence
	const float oldWeight = clamp(oldSsao.y, .0f, temporalCoherence);
	float newWeight;

	//const float oldNumSamples = oldSsao.y;
	
	if ((temporalCoherence > 1e-6f) &&
		(tex.x >= 0.0f) && (tex.x < 1.0f) && 
		(tex.y >= 0.0f) && (tex.y < 1.0f) && 
		(abs(depthDif) < MIN_DEPTH_DIFF)
		// check if something changed in the surrounding area
		//&& (oldNumSamples > 0.2 * gi.ao.y)
		)
	{
		newWeight = oldWeight + 1;

		float4 tmp = float4(gi.ao.x, gi.illum);
		float4 oldTmp = float4(oldSsao.x, oldIllum);

		float4 interpol = (tmp + oldTmp * oldWeight) / newWeight;

		OUT.ssao_col.x = interpol.x;
		OUT.illum_col.xyz = interpol.yzw;
	}
	else
	{
		newWeight = 0;

		OUT.ssao_col.x = gi.ao.x;
		OUT.illum_col.xyz = gi.illum;
	}

	OUT.ssao_col.y = newWeight;
	OUT.ssao_col.z = currentDepth;

	return OUT;
}


pixel combine(fragment IN, 
			  uniform sampler2D colorsTex,
			  uniform sampler2D ssaoTex,
			  uniform sampler2D illumTex
			  )
{
	pixel OUT;

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

	float3 illum = tex2Dlod(illumTex, float4(IN.texCoord, 0, 0)).xyz;

	OUT.illum_col.xyz = (col.xyz + illum) * ao;
	//OUT.illum_col.xyz = col.xyz * ao;

	OUT.illum_col.w = col.w;

	return OUT;
}