// Copyright (C) 2011 by The Alexandra Institute Ltd.
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#version 330
#extension GL_EXT_geometry_shader4 : enable

layout(points) in;
layout(points, max_vertices = 1) out;

precision mediump float;

// Attributes
out vec4 kInjectValueR;	// Red channel flux encoded in SH coefficients
out vec4 kInjectValueG; // Green channel flux encoded in SH coefficients
out vec4 kInjectValueB; // Blue channel flux encoded in SH coefficients

// Uniforms
uniform vec3 kLightPosition;	// World-space light position
uniform vec3 kGridOrigin;		// World-space grid origin
uniform float fInvWSGridSize;	// Inverse world-space grid cell size
uniform int iGridSize;			// Grid cell count

uniform sampler2D kWSPosTexture;		// Injection RSM input
uniform sampler2D kWSNormDepthTexture;	// Injection RSM input
uniform sampler2D kFluxTexture;			// Injection RSM input

uniform float fInvIOR;	// Inverse index of refraction

uniform int iDownsampleScale;	// Scale used to downsample RSM texture size before injection

// Constants
const float PI = 3.14159265359;

// Shader code
vec4 SHRotate(vec3 kDir, vec2 kZHCoeffs) 
{ 
	vec2 kTheta12CS = vec2(0.0);
	vec2 kPhi12CS	= vec2(0.0); 

	if (kDir.z * kDir.z < 0.99)
	{
		kTheta12CS = normalize(kDir.xy);
		kPhi12CS.x = sqrt(1.0 - kDir.z * kDir.z); 	
	}

	kPhi12CS.y = kDir.z;
	
	vec4 kResult;
	kResult.x = kZHCoeffs.x;
	kResult.y = kZHCoeffs.y * kPhi12CS.x * kTheta12CS.y; 
	kResult.z = kZHCoeffs.y * kPhi12CS.y; 
	kResult.w = kZHCoeffs.y * kPhi12CS.x * kTheta12CS.x;
	
	return kResult;
}

vec4 SHProjectCone(vec3 kDir) 
{ 
	const vec2 kZHCoeffs = vec2(0.886226925452758, 1.023326707946488); 	
	return SHRotate(kDir, kZHCoeffs);
}

void main()
{
	// Sample RSM input textures
	vec4 kWSPos			= texture2D(kWSPosTexture, gl_in[0].gl_Position.xy);  
	vec4 kWSNormDepth	= texture2D(kWSNormDepthTexture, gl_in[0].gl_Position.xy);
	vec4 kFlux			= texture2D(kFluxTexture, gl_in[0].gl_Position.xy);

	// Multiply sampled flux to compensate for downscaled sampling
	vec3 kIncomingFlux = kFlux.xyz * iDownsampleScale;

	// Reject empty cells
	if (kIncomingFlux == vec3(0.0, 0.0, 0.0))
		return;
	
	// Prepare vectors
	vec3 kNormal = normalize(kWSNormDepth.xyz);
	vec3 kLightDirection = normalize(kLightPosition - kWSPos.xyz);
	
	// Debug checks cause by malformed triangle data in mesh.
	if (isnan(kNormal.x) || isnan(kNormal.y) || isnan(kNormal.z))
		return;
	if (isinf(kNormal.x) || isinf(kNormal.y) || isinf(kNormal.z))
		return;
	if (kNormal == vec3(0.0, 0.0, 0.0))
		return;

	// Calculate fresnel
	vec3 kRefractDir = refract(-kLightDirection, kNormal, fInvIOR);
	kRefractDir = normalize(kRefractDir);
	
	float fThetaI = acos(max(0.0, dot(kNormal, kLightDirection)));
	float fThetaR = acos(max(0.0, dot(-kNormal, kRefractDir)));

	float fFresnel = 0.0;
	if (fThetaI == 0.0)
	{
		fFresnel = pow(((1.0 / fInvIOR) - 1.0) / ((1.0 / fInvIOR) + 1.0), 2.0);
	}
	else
	{
		fFresnel = 0.5 *
			(
			  (sin(fThetaI - fThetaR) * sin(fThetaI - fThetaR)) /
			  (sin(fThetaI + fThetaR) * sin(fThetaI + fThetaR))
			) *
			(
			1.0 +
			(cos(fThetaI + fThetaR) * cos(fThetaI + fThetaR)) /
			(cos(fThetaI - fThetaR) * cos(fThetaI - fThetaR))
			);
	}
	
	// Calculate flux 
	vec3 kTransmittedFlux = kIncomingFlux * (1.0 - fFresnel);
	vec4 kSH = SHProjectCone(kRefractDir);
	kInjectValueR = kTransmittedFlux.r * kSH / PI;
	kInjectValueG = kTransmittedFlux.g * kSH / PI;
	kInjectValueB = kTransmittedFlux.b * kSH / PI;
	
	// Calculate grid cell clip-space position
	vec3 kNormGridPos = (kWSPos.xyz - kGridOrigin) * fInvWSGridSize;
	vec3 kGSGridPos = kNormGridPos * iGridSize;
	
	kGSGridPos = floor(kGSGridPos + 0.5);
	kNormGridPos = kGSGridPos / iGridSize;

	vec2 kCSPos = vec2(2.0 * kNormGridPos - 1.0).xy;

	// Set inject layer and emit vertex into LPV grid
	gl_Layer = int(kGSGridPos.z);
	gl_Position = vec4(kCSPos, 0.0, 1.0);
	EmitVertex();
	EndPrimitive();
}