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

#include "../shaderenv.h"


struct fragment
{
	 // normalized screen position
	float4 pos: WPOS;
	float4 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 w,
				   uniform float3 eyePos,
				   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 ++) 
	{
		float2 offset = samples[i];

#if 1
		////////////////////
		// add random noise: 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 * AREA_SIZE * w;

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


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

		float4 sample = tex2Dlod(colors, float4(texcoord, 0, 0));
		const float eyeSpaceDepth = sample.w;
		//float3 rotView = normalize(Interpol(texcoord, bl, br, tl, tr));
		float3 rotView = Interpol(texcoord, bl, br, tl, tr);
		
		
		const float3 sample_position = eyePos - rotView * eyeSpaceDepth;
		const float3 sample_color = sample.xyz;

		float3 vector_to_sample = sample_position - centerPosition.xyz;
		const float length_to_sample = length(vector_to_sample);

		float3 direction_to_sample = vector_to_sample / length_to_sample;

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

		// distance between current position and sample position controls AO intensity.
		const float distance_intensity = 
			(SAMPLE_INTENSITY * DISTANCE_SCALE) / (DISTANCE_SCALE + length_to_sample * length_to_sample);

		// if normal perpenticular to view dir, only half of the samples count
#if 0
		const float view_correction = 1.0f + VIEW_CORRECTION_SCALE * (1.0f - dot(currentViewDir, currentNormal));
		total_color.w -= cos_angle * distance_intensity * view_correction;
		total_color.xyz += cos_angle * distance_intensity * view_correction * sample_color * ILLUM_INTENSITY;
#endif
		total_ao += cos_angle * distance_intensity;
		total_color += cos_angle * distance_intensity * sample_color * ILLUM_INTENSITY;
	}

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

	//return saturate(total_color);
	return gi;
}



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

	float4 norm = tex2D(normals, IN.texCoord.xy);
	float3 normal = normalize(norm.xyz);	

	// the w coordinate from the persp. projection
	float w = norm.w;


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

	float3 viewDir = IN.view;
	const float eyeDepth = tex2Dlod(colors, float4(IN.texCoord.xy, 0, 0)).w;
	
	float3 centerPosition;
	centerPosition.xyz = eyePos - viewDir * eyeDepth;
	
	//const float3 centerPosition = tex2D(positions, IN.texCoord.xy).xyz;


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

	GiStruct gi = globIllum(IN, colors, noiseTexture, samples, normal, centerPosition, w, eyePos, bl, br, tl, tr);
	

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

	float4 realPos = float4(centerPosition * maxDepth, 1.0f);


	// calculcate the current projected depth for next frame
	float4 currentPos = mul(modelViewProj, realPos);
	currentPos /= currentPos.w;
	const float currentDepth = currentPos.z;
	

	///////////
	//-- reprojection new frame into old one 
	
	// calculate projected depth
	float4 projPos = mul(oldModelViewProj, realPos);
	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
	float4 oldSsao = tex2D(oldSsaoTex, tex);
	float4 oldIllum = tex2D(oldIllumTex, tex);

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

	float oldWeight = clamp(oldSsao.z, 0, temporalCoherence);

	float newWeight;

	//const float oldNumSamples = oldSsao.y;
	//const float oldAvgDepth = oldSsao.z;

	if (//(temporalCoherence > 0.0f) &&
		(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;

		OUT.ssao_col.xy = (gi.ao + oldSsao.xy * oldWeight) / newWeight;
		OUT.illum_col.xyz  = (gi.illum + oldIllum.xyz * oldWeight) / newWeight;
	}
	else
	{
		newWeight = 0;

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

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

	return OUT;
}


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

	float4 col = tex2D(colors, IN.texCoord.xy);
	float ao = tex2D(ssaoTex, IN.texCoord.xy).x;
	float4 illum = tex2D(illumTex, IN.texCoord.xy);
	
	OUT.illum_col = (col + illum) * ao;
	OUT.illum_col.w = col.w;

	return OUT;
}