#version 450

#include "compiled.inc"
#include "std/math.glsl"
#include "std/gbuffer.glsl"

uniform sampler2D gbuffer0;
uniform sampler2D gbuffer1;
uniform sampler2D gbufferD;
#ifdef _EmissionShaded
uniform sampler2D gbufferEmission;
#endif

uniform mat4 P;
uniform mat4 invP;
uniform mat3 V3;

#ifdef _Sun
uniform vec3 sunDir;
uniform vec3 sunCol;
#endif

#ifdef _CPostprocess
uniform vec3 PPComp12;
#endif

in vec2 texCoord;
out vec4 fragColor;

const float GOLDEN_ANGLE = 2.39996323;
const int RAY_STEPS = 6;
const int BINARY_STEPS = 3;

vec2 getProjectedCoord(const vec3 viewPos) {
	vec4 projectedCoord = P * vec4(viewPos, 1.0);
	if (projectedCoord.w <= 0.0) return vec2(-1e5);
	projectedCoord.xy /= projectedCoord.w;
	projectedCoord.xy = projectedCoord.xy * 0.5 + 0.5;
	#ifdef _InvY
	projectedCoord.y = 1.0 - projectedCoord.y;
	#endif
	return projectedCoord.xy;
}

float linearZ(const float depth) {
	return -P[3].z / (depth + P[2].z);
}

void buildTBN(vec3 n, out vec3 tangent, out vec3 bitangent) {
	vec3 absN = abs(n);
	if (absN.x <= absN.y && absN.x <= absN.z) {
		tangent = normalize(cross(n, vec3(1.0, 0.0, 0.0)));
	} else if (absN.y <= absN.z) {
		tangent = normalize(cross(n, vec3(0.0, 1.0, 0.0)));
	} else {
		tangent = normalize(cross(n, vec3(0.0, 0.0, 1.0)));
	}
	bitangent = cross(n, tangent);
}

vec3 sampleHemisphere(vec3 n, vec3 tangent, vec3 bitangent, float phi, float cosTheta) {
	float sinTheta = sqrt(1.0 - cosTheta * cosTheta);
	vec3 h = vec3(cos(phi) * sinTheta, sin(phi) * sinTheta, cosTheta);
	return normalize(tangent * h.x + bitangent * h.y + n * h.z);
}

vec3 traceRay(vec3 origin, vec3 dir, float maxDist, float minDist, float jitter) {
	float stepSize = maxDist / float(RAY_STEPS);
	float rayDist = minDist + stepSize * (jitter - 1.0);
	vec3 pos = origin + dir * rayDist;
	
	float prevDepthDiff = 0.0;
	bool hasPrev = false;
	
	for (int i = 0; i < RAY_STEPS; i++) {
		pos += dir * stepSize;
		rayDist += stepSize;
		vec2 uv = getProjectedCoord(pos);
		if (uv.x < -100.0) return vec3(-1.0);
		uv = clamp(uv, vec2(0.001), vec2(0.999));
		
		float sampleDepth = textureLod(gbufferD, uv, 0.0).r * 2.0 - 1.0;
		if (sampleDepth == 1.0) {
			hasPrev = false;
			continue;
		}

		float depthDiff = pos.z - linearZ(sampleDepth);
		float thickness = maxDist * 0.075 + rayDist * 0.125;
		
		bool crossed = hasPrev && (prevDepthDiff > 0.0) && (depthDiff <= 0.0);
		bool withinThickness = (depthDiff <= 0.0) && (-depthDiff < thickness);
		
		if (crossed || withinThickness) {
			vec3 bPos = pos;
			vec3 bDir = dir * stepSize;
			for (int j = 0; j < BINARY_STEPS; j++) {
				bDir *= 0.5;
				bPos -= bDir;
				vec2 bUV = getProjectedCoord(bPos);
				bUV = clamp(bUV, vec2(0.001), vec2(0.999));
				float bDepth = textureLod(gbufferD, bUV, 0.0).r * 2.0 - 1.0;
				if (bDepth == 1.0) {
					bPos += bDir;
					continue;
				}
				if (bPos.z - linearZ(bDepth) > 0.0) bPos += bDir;
			}
			
			vec2 bestUV = getProjectedCoord(bPos);
			if (bestUV.x < -100.0) return vec3(-1.0);
			bestUV = clamp(bestUV, vec2(0.001), vec2(0.999));
			
			float hitDist = length(bPos - origin);
			float distWeight = max(0.0, 1.0 - (hitDist / maxDist));
			distWeight *= distWeight;
			
			return vec3(bestUV, distWeight);
		}
		
		prevDepthDiff = depthDiff;
		hasPrev = true;
	}
	
	return vec3(-1.0);
}

void main() {
	float depth = textureLod(gbufferD, texCoord, 0.0).r * 2.0 - 1.0;
	if (depth == 1.0) {
		fragColor = vec4(0.0, 0.0, 0.0, 1.0);
		return;
	}

	vec4 g0 = textureLod(gbuffer0, texCoord, 0.0);
	vec3 n = getNor(g0.rg);
	vec3 basecolor = textureLod(gbuffer1, texCoord, 0.0).rgb;

	vec3 viewNormal = V3 * n;
	vec3 viewPos = getPosView2(invP, depth, texCoord);

	#ifdef _CPostprocess
		float radius = PPComp12.y * 2.0;
		float strength = PPComp12.x * 0.5;
	#else
		float radius = ssgiRadius * 2.0;
		float strength = ssgiStrength * 0.5;
	#endif

	radius = min(radius, max(-viewPos.z / P[1].y, 0.05));

	float noise = fract(52.9829189 * fract(0.06711056 * gl_FragCoord.x + 0.00583715 * gl_FragCoord.y));

	vec3 gi = vec3(0.0);
	float missWeight = 0.0;

	// min distance to avoid self shadowing artiffacts
	float minDist = radius * 0.05;
	
	vec3 tangent, bitangent;
	buildTBN(viewNormal, tangent, bitangent);
	
	for (int i = 0; i < ssgiSamples; i++) {
		float phi = float(i) * GOLDEN_ANGLE + noise * PI2;
		float cosTheta = sqrt(max(0.0, 1.0 - (float(i) + 0.5) / float(ssgiSamples)));
		float jitter = fract(noise + float(i) * 0.618034);
		
		vec3 rayDir = sampleHemisphere(viewNormal, tangent, bitangent, phi, cosTheta);
		vec3 hitResult = traceRay(viewPos, rayDir, radius, minDist, jitter);
		
		if (hitResult.x < 0.0) {
			missWeight += 1.0;
			continue;
		}
		
		vec2 hitUV = hitResult.xy;
		float distWeight = hitResult.z;
		
		vec3 hitN = getNor(textureLod(gbuffer0, hitUV, 0.0).rg);
		float emitterCos = max(0.0, dot(V3 * hitN, -rayDir));
		if (emitterCos <= 0.0) {
			missWeight += 1.0;
			continue;
		}
		
		vec3 hitAlbedo = textureLod(gbuffer1, hitUV, 0.0).rgb;
		
		#ifdef _Sun
		float hitNdotL = max(0.0, dot(hitN, sunDir));
		vec3 hitRadiance = hitAlbedo * (sunCol * hitNdotL + vec3(0.4));
		#else
		vec3 hitRadiance = hitAlbedo * 0.4;
		#endif
		
		#ifdef _EmissionShaded
		hitRadiance += textureLod(gbufferEmission, hitUV, 0.0).rgb;
		#endif
		
		gi += hitRadiance * (distWeight * emitterCos);
	}

	gi += basecolor * 0.1 * missWeight;
	gi *= 2.0 * strength / float(ssgiSamples);

	#ifdef _EmissionShaded
	gi += textureLod(gbufferEmission, texCoord, 0.0).rgb * 0.3;
	#endif

	fragColor = vec4(min(gi, vec3(2.0)), 1.0);
}