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


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 positions,
				   uniform sampler2D noiseTexture,
				   uniform float2 samples[NUM_SAMPLES],
				   uniform float3 currentNormal,
				   uniform float4 centerPosition,
				   float w,
				   // uniform float3 viewDir,
				   uniform float3 eyePos,
				   uniform float3 bl,
				   uniform float3 br,
				   uniform float3 tl,
				   uniform float3 tr
				   )
{
	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, SSAO_MIPMAP_LEVEL));
		const float eyeSpaceDepth = sample.w;
		float3 rotView = normalize(Interpol(texcoord, bl, br, tl, tr));
		
		const float3 sample_position = eyePos - rotView * eyeSpaceDepth;
		const float3 sample_color = sample.xyz;

		// use lower lod level to improve cache coherence
		//float3 sample_position = tex2Dlod(positions, float4(texcoord, 0, SSAO_MIPMAP_LEVEL)).xyz;
		//float3 sample_color = tex2Dlod(colors, float4(texcoord, 0, GI_MIPMAP_LEVEL)).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;

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


/** The mrt shader for screen space ambient occlusion + indirect illumination
*/
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,
		   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);	
	// something like a constant ambient term
	const float amb = norm.w;
	
	/// the current view direction
	//float3 viewDir = normalize(IN.view);

	// the w coordinate from the persp. projection
	float w = norm.w;
	// the current world position
	const float4 centerPosition = tex2D(positions, IN.texCoord.xy);
	// the current color
	const float4 currentCol = tex2Dlod(colors, float4(IN.texCoord.xy, 0, 0));
	// the current depth is stored in the w component
	const float currentDepth = centerPosition.w;

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

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

	float4 realPos = centerPosition * maxDepth;
	realPos.w = 1.0f;

	float4 oldPos = mul(oldModelViewProj, realPos);

	const float newDepth = oldPos.z / oldPos.w;

	float2 tex = (oldPos.xy / oldPos.w) * 0.5f + 0.5f;

	float4 oldSsao = tex2D(oldSsaoTex, tex);
	float4 oldIllum = tex2D(oldIllumTex, tex);

	const float oldDepth = oldSsao.w;
	const float depthDif = 1.0f - newDepth / 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) < 1e-3f)
		// check if something changed in the surrounding area
		&& (oldNumSamples > 0.2 * gi.ao.y)
		//&& (oldAvgDepth / newAvgDepth > 0.99)
		)
	{
		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;
}