////////////////////
// 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, 
				   float3 viewDir,
				   float sampleIntensity
				   )
{
	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(sampleIntensity, 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,
		   uniform float kernelRadius,
		   uniform float sampleIntensity
		   )
{
	pixel2 OUT;

	const float3 normal = tex2Dlod(normals, float4(IN.texCoord, 0 ,0)).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 = kernelRadius / currentPos.w;
	currentPos /= currentPos.w;
	
	const float currentDepth = currentPos.z;

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

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

	/////////////////
	//-- 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;

	// 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;
}