#include "../shaderenv.h"

////////////////////
// Screen Spaced Ambient Occlusion shader
// based on shader of Alexander Kusternig


#define USE_EYE_SPACE_DEPTH 1


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


struct pixel
{
	float4 illum_col: COLOR0;
};


inline float occlusionPower(float radius, float dist)
{
	return 6.283185307179586476925286766559f * (1.0f - cos(asin(radius / dist)));
}


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


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


// reconstruct world space position
inline float3 ReconstructSamplePos(uniform sampler2D tex,
								   float2 texcoord, 
								   float3 bl, float3 br, float3 tl, float3 tr)
{
	const float eyeSpaceDepth = tex2D(tex, texcoord).w;
	float3 viewVec = Interpol(texcoord, bl, br, tl, tr);
	float3 samplePos = -viewVec * eyeSpaceDepth;

	return samplePos;
}

#pragma position_invariant temporalSmoothing

/** This shader computes the reprojection and stores reprojected color / depth values
	as well as a boolean that 
*/
float4 temporalSmoothing(float4 currentProjPos,
						 float4 worldPos,
						 float currentDepth,
						 uniform sampler2D oldTex,
						 const uniform float4x4 oldModelViewProj,
						 uniform float temporalCoherence,
						 uniform float2 ao,
						 uniform float2 samples[NUM_SAMPLES],
						 uniform sampler2D colors,
						 uniform sampler2D noiseTex,
						 uniform float scaleFactor,
						 uniform float3 bl,
						 uniform float3 br,
						 uniform float3 tl,
						 uniform float3 tr, 
						 float2 texcoord0,
						 float3 eyePos,
						 float eyeSpaceDepth,
						 float3 oldEyePos,
						 uniform float3 oldbl,
						 uniform float3 oldbr,
						 uniform float3 oldtl,
						 uniform float3 oldtr
						 )
{
	float4 illum_col;

	/////////////////
	//-- compute reprojection for temporal smoothing

	// reproject into old frame and calculate projected depth
	float4 backProjPos = mul(oldModelViewProj, worldPos);
	backProjPos /= backProjPos.w;
	// fit from unit cube into 0 .. 1
	const float2 oldTexCoords = backProjPos.xy * 0.5f + 0.5f;

	// retrieve the sample from the last frame
	float4 oldCol = tex2D(oldTex, oldTexCoords);
	const float oldEyeSpaceDepth = oldCol.z;

	float3 viewVec = Interpol(oldTexCoords, oldbl, oldbr, oldtl, oldtr);
	//float3 oldSamplePos = -viewVec * eyeSpaceDepth + oldEyePos;
	float3 oldSamplePos = - viewVec * oldCol.z;
	oldSamplePos += oldEyePos;

	//const float oldDepth = oldCol.z;
	const float depthDif = length(oldSamplePos - worldPos.xyz);

	//const float oldNumSamples = oldCol.y;
	const float oldWeight = clamp(oldCol.y, .0f, temporalCoherence);
	float newWeight;

	/*bool isValid = true;

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

		float2 mynoise = tex2D(noiseTex, texcoord0).xy;
		const float2 offsetTransformed = myreflect(offset, mynoise);

		const float2 texCoord = texcoord0;// + offsetTransformed * scaleFactor;
		const float3 samplePos = ReconstructSamplePos(colors, texCoord, bl, br, tl, tr) + eyePos;

		// reproject into old frame and calculate projected depth
		float4 projPos = mul(oldModelViewProj, float4(samplePos, 1.0f));
		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
		// retrieve the sample from the last frame
		const float4 oldSample = tex2D(oldTex, projPos.xy * 0.5f + 0.5f);
		const float dDiff = projDepth - oldSample.z;

		if (abs(dDiff) > 1e-5f) isValid = false;
	}*/

	// the number of valid samples in this frame
	//const float newNumSamples = ao.y;

	if ((temporalCoherence > 1e-6f)
		&& (oldTexCoords.x >= 0.0f) && (oldTexCoords.x < 1.0f)
		&& (oldTexCoords.y >= 0.0f) && (oldTexCoords.y < 1.0f)
		&& (abs(depthDif) <= MIN_DEPTH_DIFF) 
		// if visibility changed in the surrounding area we have to recompute
		//&& (oldNumSamples > 0.8f * newNumSamples)
	//	&& isValid
		)
	{
		// increase the weight for convergence
		newWeight = oldWeight + 1.0f;
		illum_col.x = (ao.x + oldCol.x * oldWeight) / newWeight;
		//if (!(oldNumSamples > ao.y - 1.5f)) newWeight = 0;
	}
	else
	{	
		illum_col.x = ao.x;
		newWeight = .0f;
	}
	
	//isValid = 0.0f;
	//illum_col.y = isValid / 16.0f;
	illum_col.y = newWeight;
	illum_col.z = eyeSpaceDepth;

	return illum_col;
}


/** The ssao shader returning the an intensity value between 0 and 1
*/
float2 ssao(fragment IN,
			uniform sampler2D colors,
			uniform sampler2D noiseTex,
			uniform float2 samples[NUM_SAMPLES],
			uniform float3 currentNormal,
			uniform float3 centerPosition,
			uniform float scaleFactor,
			uniform float3 bl,
			uniform float3 br,
			uniform float3 tl,
			uniform float3 tr, 
			uniform float3 viewDir
			)
{
	// 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.

	float total_ao = .0f;
	float numSamples = .0f;


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

#if 1
		////////////////////
		//-- add random noise: reflect around random normal vector (rather slow!)

		float2 mynoise = tex2D(noiseTex, IN.texCoord).xy;
		const float2 offsetTransformed = myreflect(offset, mynoise);
#else
		const float2 offsetTransformed = offset;
#endif
		// weight with projected coordinate to reach similar kernel size for near and far
		const 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;

		const float3 samplePos = ReconstructSamplePos(colors, texcoord, bl, br, tl, tr);


		////////////////
		//-- compute contribution of sample using the direction and angle

		float3 dirSample = samplePos - centerPosition;
		const float lengthToSample = length(dirSample);
		// normalize
		dirSample /= lengthToSample;

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

		// the distance_scale offset is used to avoid singularity that occurs at global illumination when 
		// the distance to a sample approaches zero
		//const float aoContrib = (1.0f > lengthToSample) ? occlusionPower(9e-2f, DISTANCE_SCALE + lengthToSample): .0f;
		const float aoContrib = SAMPLE_INTENSITY / (DISTANCE_SCALE + lengthToSample * lengthToSample);

#if 1
		// if surface normal perpenticular to view dir, approx. half of the samples will not count
		// => compensate for this (on the other hand, projected sampling area could be larger!)
		const float viewCorrection = 1.0f + VIEW_CORRECTION_SCALE * dot(viewDir, currentNormal);
		total_ao += cosAngle * aoContrib * viewCorrection;
#else
		total_ao += cosAngle * intensity;
#endif
	}

	return float2(max(0.0f, 1.0f - total_ao), numSamples);
}

//#pragma position_invariant main

/** The mrt shader for screen space ambient occlusion
*/
pixel main(fragment IN, 
		   uniform sampler2D colors,
		   uniform sampler2D normals,
		   uniform sampler2D noiseTex,
		   uniform float2 samples[NUM_SAMPLES],
		   uniform sampler2D oldTex,
		   uniform float4x4 modelViewProj,
		   uniform float4x4 oldModelViewProj,
		   uniform float temporalCoherence,
		   uniform float3 eyePos,
		   uniform float3 bl,
		   uniform float3 br,
		   uniform float3 tl,
		   uniform float3 tr,
		   uniform float3 oldEyePos,
		   uniform float3 oldbl,
		   uniform float3 oldbr,
		   uniform float3 oldtl,
		   uniform float3 oldtr
		   )
{
	pixel OUT;

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

	// reconstruct position from the eye space depth
	const float3 viewDir = IN.view;
	const float eyeDepth = tex2Dlod(colors, float4(IN.texCoord, 0, 0)).w;
	const float3 eyeSpacePos = -viewDir * eyeDepth;
	const float4 worldPos = float4(eyePos + eyeSpacePos, 1.0f);


	////////////////
	//-- calculcate the current projected posiion (also used for next frame)
	
	float4 projPos = mul(modelViewProj, worldPos);
	
	float w = 1.0f / projPos.w;
	projPos *= w;
	w *= SAMPLE_RADIUS;
	
	const float currentDepth = projPos.z * PRECISION_SCALE;

	const float2 ao = ssao(IN, colors, noiseTex, samples, normal, eyeSpacePos, w, bl, br, tl, tr, normalize(viewDir));


	/////////////////
	//-- compute temporally smoothing
	
	OUT.illum_col = temporalSmoothing(projPos, worldPos, currentDepth, oldTex, oldModelViewProj, temporalCoherence, ao,
		                              samples, colors, noiseTex, w, bl, br, tl, tr, IN.texCoord, eyePos, eyeDepth,
									  oldEyePos, oldbl, oldbr, oldtl, oldtr);

	return OUT;
}


float Filter(float2 texCoord, 
			 uniform sampler2D ssaoTex,
			 uniform float2 filterOffs[NUM_SSAO_FILTERSAMPLES],
			 uniform float filterWeights[NUM_SSAO_FILTERSAMPLES]
)
{
	float average = .0f;
	float w = .0f;

	for (int i = 0; i < NUM_SSAO_FILTERSAMPLES; ++ i)
	{
		average += filterWeights[i] * tex2Dlod(ssaoTex, float4(texCoord + filterOffs[i], 0, 0)).x;
		w += filterWeights[i];
	}

	average *= 1.0f / (float)w;

	return average;
}


pixel combine(fragment IN, 
			  uniform sampler2D colors,
			  uniform sampler2D ssaoTex,
			  uniform float2 filterOffs[NUM_SSAO_FILTERSAMPLES],
			  uniform float filterWeights[NUM_SSAO_FILTERSAMPLES]
			  )
{
	pixel OUT;

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

	//if (ao.y < 10.0f) 
	//	ao.x = Filter(IN.texCoord, ssaoTex, filterOffs, filterWeights);

	//OUT.illum_col = col * ao.x;
	OUT.illum_col.xyz = float3(1.0f - ao.x, 1.0f - ao.y * 1e-2f, 1);
	//OUT.illum_col.xyz = float3(1.0f - ao.x, ao.y, 0);
	OUT.illum_col.w = col.w;

	return OUT;
}