Handling multiple light sources (HDRI, Quad light, Sphere light)

HDRI_Environment_Fragment.glsl

precision highp float;
precision highp int;
precision highp sampler2D;

uniform sampler2D tTriangleTexture;
uniform sampler2D tAABBTexture;
uniform sampler2D tHDRTexture;

#include <pathtracing_uniforms_and_defines>

uniform vec3 uMaterialColor;
uniform vec3 uHDRIColor;
uniform vec3 uSunDirectionVector;
uniform float uHDRI_Exposure;
uniform int uUSE_HDRI;
uniform float uRoughness;
uniform int uMaterialType;

//float InvTextureWidth = 0.000244140625; // (1 / 4096 texture width)
//float InvTextureWidth = 0.00048828125;  // (1 / 2048 texture width)
//float InvTextureWidth = 0.0009765625;   // (1 / 1024 texture width)

#define INV_TEXTURE_WIDTH 0.00048828125

// the following directions pointing at the sun were found by trial and error: left here just for reference
//#define SUN_DIRECTION normalize(vec3(-0.555, 1.0, 0.205)) // use this vec3 for the symmetrical_garden_2k.hdr environment
//#define SUN_DIRECTION normalize(vec3(0.54, 1.0, -0.595)) // use this vec3 for the kiara_5_noon_2k.hdr environment
 
#define N_SPHERES 1
#define N_QUADS 1
// #define N_BOXES 2

//-----------------------------------------------------------------------

vec3 rayOrigin, rayDirection;
// recorded intersection data:
vec3 hitNormal, hitEmission, hitColor;
vec2 hitUV;
int hitType;
float hitObjectID;
bool hitIsModel;

struct Quad { vec3 normal; vec3 v0; vec3 v1; vec3 v2; vec3 v3; vec3 emission; vec3 color; int type; };
struct Sphere { float radius; vec3 position; vec3 emission; vec3 color; int type; };
// struct Box { vec3 minCorner; vec3 maxCorner; vec3 emission; vec3 color; int type; };

Sphere spheres[N_SPHERES];
Quad quads[N_QUADS];
// Box boxes[N_BOXES];


#include <pathtracing_random_functions>

#include <pathtracing_calc_fresnel_reflectance>

#include <pathtracing_sphere_intersect>

// #include <pathtracing_box_intersect>

#include <pathtracing_boundingbox_intersect>

#include <pathtracing_bvhTriangle_intersect>


#include <pathtracing_quad_intersect>

//#include <pathtracing_bvhDoubleSidedTriangle_intersect>

#include <pathtracing_sample_sphere_light>	

#include <pathtracing_sample_quad_light>

vec2 stackLevels[28];

//vec4 boxNodeData0 corresponds to .x = idTriangle,  .y = aabbMin.x, .z = aabbMin.y, .w = aabbMin.z
//vec4 boxNodeData1 corresponds to .x = idRightChild .y = aabbMax.x, .z = aabbMax.y, .w = aabbMax.z

void GetBoxNodeData(const in float i, inout vec4 boxNodeData0, inout vec4 boxNodeData1)
{
	// each bounding box's data is encoded in 2 rgba(or xyzw) texture slots 
	float ix2 = i * 2.0;
	// (ix2 + 0.0) corresponds to .x = idTriangle,  .y = aabbMin.x, .z = aabbMin.y, .w = aabbMin.z 
	// (ix2 + 1.0) corresponds to .x = idRightChild .y = aabbMax.x, .z = aabbMax.y, .w = aabbMax.z 

	ivec2 uv0 = ivec2( mod(ix2 + 0.0, 2048.0), (ix2 + 0.0) * INV_TEXTURE_WIDTH ); // data0
	ivec2 uv1 = ivec2( mod(ix2 + 1.0, 2048.0), (ix2 + 1.0) * INV_TEXTURE_WIDTH ); // data1
	
	boxNodeData0 = texelFetch(tAABBTexture, uv0, 0);
	boxNodeData1 = texelFetch(tAABBTexture, uv1, 0);
}


//-----------------------------------------------------------------------------------------------------------------------------------------------
float SceneIntersect( out bool finalIsRayExiting )
//-----------------------------------------------------------------------------------------------------------------------------------------------
{
        vec4 currentBoxNodeData0, nodeAData0, nodeBData0, tmpNodeData0;
	vec4 currentBoxNodeData1, nodeAData1, nodeBData1, tmpNodeData1;
	
	vec4 vd0, vd1, vd2, vd3, vd4, vd5, vd6, vd7;

	vec3 inverseDir = 1.0 / rayDirection;
	vec3 normal;

	vec2 currentStackData, stackDataA, stackDataB, tmpStackData;
	ivec2 uv0, uv1, uv2, uv3, uv4, uv5, uv6, uv7;

	float d;
	float t = INFINITY;
        float stackptr = 0.0;
	float id = 0.0;
	float tu, tv;
	float triangleID = 0.0;
	float triangleU = 0.0;
	float triangleV = 0.0;
	float triangleW = 0.0;

	int objectCount = 0;
	
	hitObjectID = -INFINITY;

	bool skip = false;
	bool triangleLookupNeeded = false;
	bool isRayExiting = false;
	hitIsModel = false;

	
	for (int i = 0; i < N_SPHERES; i++)
        {
		d = SphereIntersect( spheres[i].radius, spheres[i].position, rayOrigin, rayDirection );
		if (d < t)
		{
			t = d;
			hitNormal = (rayOrigin + rayDirection * t) - spheres[i].position;
			hitEmission = spheres[i].emission;
			hitColor = spheres[i].color;
			hitType = spheres[i].type;
			hitIsModel = false;
			hitObjectID = float(objectCount);
		}
		objectCount++;
	}

	for (int i = 0; i < N_QUADS; i++)
        {
		d = QuadIntersect( quads[i].v0, quads[i].v1, quads[i].v2, quads[i].v3, rayOrigin, rayDirection, false );
		if (d < t)
		{
			t = d;
			hitNormal = quads[i].normal;
			hitEmission = quads[i].emission;
			hitColor = quads[i].color;
			hitType = quads[i].type;
			hitObjectID = float(objectCount);
		}
		objectCount++;
        }
		
	// for (int i = 0; i < N_BOXES; i++)
    //     {
	// 	d = BoxIntersect( boxes[i].minCorner, boxes[i].maxCorner, rayOrigin, rayDirection, normal, isRayExiting );
	// 	if (d < t)
	// 	{
	// 		t = d;
	// 		hitNormal = normal;
	// 		hitEmission = boxes[i].emission;
	// 		hitColor = boxes[i].color;
	// 		hitType = boxes[i].type;
	// 		hitIsModel = false;
	// 		finalIsRayExiting = isRayExiting;
	// 		hitObjectID = float(objectCount);
	// 	}
	// 	objectCount++;
	// }



	GetBoxNodeData(stackptr, currentBoxNodeData0, currentBoxNodeData1);
	currentStackData = vec2(stackptr, BoundingBoxIntersect(currentBoxNodeData0.yzw, currentBoxNodeData1.yzw, rayOrigin, inverseDir));
	stackLevels[0] = currentStackData;
	skip = (currentStackData.y < t);

	while (true)
        {
		if (!skip) 
                {
                        // decrease pointer by 1 (0.0 is root level, 27.0 is maximum depth)
                        if (--stackptr < 0.0) // went past the root level, terminate loop
                                break;

                        currentStackData = stackLevels[int(stackptr)];
			
			if (currentStackData.y >= t)
				continue;
			
			GetBoxNodeData(currentStackData.x, currentBoxNodeData0, currentBoxNodeData1);
                }
		skip = false; // reset skip
		

		if (currentBoxNodeData0.x < 0.0) // < 0.0 signifies an inner node
		{
			GetBoxNodeData(currentStackData.x + 1.0, nodeAData0, nodeAData1);
			GetBoxNodeData(currentBoxNodeData1.x, nodeBData0, nodeBData1);
			stackDataA = vec2(currentStackData.x + 1.0, BoundingBoxIntersect(nodeAData0.yzw, nodeAData1.yzw, rayOrigin, inverseDir));
			stackDataB = vec2(currentBoxNodeData1.x, BoundingBoxIntersect(nodeBData0.yzw, nodeBData1.yzw, rayOrigin, inverseDir));
			
			// first sort the branch node data so that 'a' is the smallest
			if (stackDataB.y < stackDataA.y)
			{
				tmpStackData = stackDataB;
				stackDataB = stackDataA;
				stackDataA = tmpStackData;

				tmpNodeData0 = nodeBData0;   tmpNodeData1 = nodeBData1;
				nodeBData0   = nodeAData0;   nodeBData1   = nodeAData1;
				nodeAData0   = tmpNodeData0; nodeAData1   = tmpNodeData1;
			} // branch 'b' now has the larger rayT value of 'a' and 'b'

			if (stackDataB.y < t) // see if branch 'b' (the larger rayT) needs to be processed
			{
				currentStackData = stackDataB;
				currentBoxNodeData0 = nodeBData0;
				currentBoxNodeData1 = nodeBData1;
				skip = true; // this will prevent the stackptr from decreasing by 1
			}
			if (stackDataA.y < t) // see if branch 'a' (the smaller rayT) needs to be processed 
			{
				if (skip) // if larger branch 'b' needed to be processed also,
					stackLevels[int(stackptr++)] = stackDataB; // cue larger branch 'b' for future round
							// also, increase pointer by 1
				
				currentStackData = stackDataA;
				currentBoxNodeData0 = nodeAData0; 
				currentBoxNodeData1 = nodeAData1;
				skip = true; // this will prevent the stackptr from decreasing by 1
			}

			continue;
		} // end if (currentBoxNodeData0.x < 0.0) // inner node


		// else this is a leaf

		// each triangle's data is encoded in 8 rgba(or xyzw) texture slots
		id = 8.0 * currentBoxNodeData0.x;

		uv0 = ivec2( mod(id + 0.0, 2048.0), (id + 0.0) * INV_TEXTURE_WIDTH );
		uv1 = ivec2( mod(id + 1.0, 2048.0), (id + 1.0) * INV_TEXTURE_WIDTH );
		uv2 = ivec2( mod(id + 2.0, 2048.0), (id + 2.0) * INV_TEXTURE_WIDTH );
		
		vd0 = texelFetch(tTriangleTexture, uv0, 0);
		vd1 = texelFetch(tTriangleTexture, uv1, 0);
		vd2 = texelFetch(tTriangleTexture, uv2, 0);

		d = BVH_TriangleIntersect( vec3(vd0.xyz), vec3(vd0.w, vd1.xy), vec3(vd1.zw, vd2.x), rayOrigin, rayDirection, tu, tv );

		if (d < t)
		{
			t = d;
			triangleID = id;
			triangleU = tu;
			triangleV = tv;
			triangleLookupNeeded = true;
		}
	      
        } // end while (true)
	

        if (triangleLookupNeeded)
	{
		//uv0 = ivec2( mod(triangleID + 0.0, 2048.0), (triangleID + 0.0) * INV_TEXTURE_WIDTH );
		//uv1 = ivec2( mod(triangleID + 1.0, 2048.0), (triangleID + 1.0) * INV_TEXTURE_WIDTH );
		uv2 = ivec2( mod(triangleID + 2.0, 2048.0), (triangleID + 2.0) * INV_TEXTURE_WIDTH );
		uv3 = ivec2( mod(triangleID + 3.0, 2048.0), (triangleID + 3.0) * INV_TEXTURE_WIDTH );
		uv4 = ivec2( mod(triangleID + 4.0, 2048.0), (triangleID + 4.0) * INV_TEXTURE_WIDTH );
		uv5 = ivec2( mod(triangleID + 5.0, 2048.0), (triangleID + 5.0) * INV_TEXTURE_WIDTH );
		//uv6 = ivec2( mod(triangleID + 6.0, 2048.0), (triangleID + 6.0) * INV_TEXTURE_WIDTH );
		//uv7 = ivec2( mod(triangleID + 7.0, 2048.0), (triangleID + 7.0) * INV_TEXTURE_WIDTH );
		
		//vd0 = texelFetch(tTriangleTexture, uv0, 0);
		//vd1 = texelFetch(tTriangleTexture, uv1, 0);
		vd2 = texelFetch(tTriangleTexture, uv2, 0);
		vd3 = texelFetch(tTriangleTexture, uv3, 0);
		vd4 = texelFetch(tTriangleTexture, uv4, 0);
		vd5 = texelFetch(tTriangleTexture, uv5, 0);
		//vd6 = texelFetch(tTriangleTexture, uv6, 0);
		//vd7 = texelFetch(tTriangleTexture, uv7, 0);

	
		// face normal for flat-shaded polygon look
		//hitNormal = normalize( cross(vec3(vd0.w, vd1.xy) - vec3(vd0.xyz), vec3(vd1.zw, vd2.x) - vec3(vd0.xyz)) );
		
		// interpolated normal using triangle intersection's uv's
		triangleW = 1.0 - triangleU - triangleV;
		hitNormal = triangleW * vec3(vd2.yzw) + triangleU * vec3(vd3.xyz) + triangleV * vec3(vd3.w, vd4.xy);
		hitEmission = vec3(0);
		hitColor = uMaterialColor;//vd6.yzw;
		hitUV = triangleW * vec2(vd4.zw) + triangleU * vec2(vd5.xy) + triangleV * vec2(vd5.zw); 
		hitType = int(uMaterialType);//int(vd6.x);
		//hitAlbedoTextureID = -1;//int(vd7.x);
		hitIsModel = true;
		hitObjectID = float(objectCount);
	}

	return t;

} // end float SceneIntersect( out bool finalIsRayExiting )



vec3 Get_HDR_Color(vec3 rDirection)
{
	vec2 sampleUV;
	//sampleUV.y = asin(clamp(rDirection.y, -1.0, 1.0)) * ONE_OVER_PI + 0.5;
	///sampleUV.x = (1.0 + atan(rDirection.x, -rDirection.z) * ONE_OVER_PI) * 0.5;
	sampleUV.x = atan(rDirection.x, -rDirection.z) * ONE_OVER_TWO_PI + 0.5;
  	sampleUV.y = acos(rDirection.y) * ONE_OVER_PI;
	vec3 texColor = texture(tHDRTexture, sampleUV).rgb;

	// texColor = texData.a > 0.57 ? vec3(100) : vec3(0);
	// return texColor;
	return texColor * uHDRI_Exposure;
}


//----------------------------------------------------------------------------------------------------------------------------------------------------
vec3 CalculateRadiance( out vec3 objectNormal, out vec3 objectColor, out float objectID, out float pixelSharpness )
//----------------------------------------------------------------------------------------------------------------------------------------------------
{
	vec3 accumCol = vec3(0);
        vec3 mask = vec3(1);
	vec3 checkCol0 = vec3(1);
	vec3 checkCol1 = vec3(0.5);
        vec3 tdir;
	vec3 x, n, nl;
        
	float t;
        float nc, nt, ratioIoR, Re, Tr;
	float P, RP, TP;
	float weight;
	float thickness = 0.1;
	float roughness = 0.0;
	float randomLightChoice = 0.0;

	int diffuseCount = 0;
	int previousIntersecType = -100;
	hitType = -100;

	bool coatTypeIntersected = false;
	bool bounceIsSpecular = true;
	bool sampleLight = false;
	bool isRayExiting = false;

	vec3 dirToLight;
  	//vec3 dirToLight2;
  

	for (int bounces = 0; bounces < 5; bounces++)
	{
		previousIntersecType = hitType;

		t = SceneIntersect(isRayExiting);
		roughness = hitIsModel ? uRoughness : roughness;
		
		if (hitType == LIGHT)
		{	
			if (bounceIsSpecular || sampleLight)
				accumCol = mask * hitEmission;
			// reached a light, so we can exit
			break;

		} // end if (hitType == LIGHT)

		if (t == INFINITY)
		{	
                        vec3 environmentCol = Get_HDR_Color(rayDirection);

			// looking directly at sky
			if (bounces == 0)
			{
				pixelSharpness = 1.01;
				// if (uUSE_HDRI == 0) 
				// {
				// 	// accumCol = vec3(0.001,0.001,0.001);
				// 	accumCol = uHDRIColor;
				// }
		 		// else 
				// {
				// 	accumCol = environmentCol;
				// }
			
				accumCol = environmentCol;
				break;
			}

			// sun light source location in HDRI image is sampled by a diffuse surface
			// mask has already been down-weighted in this case
			if (sampleLight)
			{
				accumCol = mask * environmentCol;
				break;
			}

			// sky image seen in a reflection or refraction
			// if (bounceIsSpecular)
			// {
			// 	// try to get rid of fireflies on rougher surfaces
			// 	if (dot(rayDirection, uSunDirectionVector) > 0.98)
			// 		environmentCol = vec3(1);
			// 		//environmentCol = mix( vec3(1), environmentCol, clamp(pow(roughness, 0.0), 0.0, 1.0) );
				
			// 	accumCol = mask * environmentCol;
			// 	break;
			// }

			// random diffuse bounce hits sky
			if ( !bounceIsSpecular )
			{
				weight = dot(rayDirection, uSunDirectionVector) < 0.98 ? 1.0 : 0.0;
				accumCol = mask * environmentCol * weight;

				// note: you don't need this line because the glass table and checker floor (huge sphere below) both are not in your scene setup
				//if (bounces == 3) accumCol = mask * environmentCol * weight * 2.0; // checkered ground beneath glass table

				break;
			}
			
			if (bounceIsSpecular)
			{
				if (coatTypeIntersected)
				{
					if (dot(rayDirection, uSunDirectionVector) > 0.998)
					pixelSharpness = 1.01;
				}
				else
					pixelSharpness = 1.01;

				
				if (dot(rayDirection, uSunDirectionVector) > 0.8)
				{
					environmentCol = mix(vec3(1), environmentCol, clamp(pow(1.0-roughness, 20.0), 0.0, 1.0));
				}
					
				accumCol = mask * environmentCol;
				break;
			}
					
			// reached the HDRI sky light, so we can exit
			break;

		} // end if (t == INFINITY)


		// useful data 
		n = normalize(hitNormal);
                nl = dot(n, rayDirection) < 0.0 ? n : -n;
		x = rayOrigin + rayDirection * t;

		if (bounces == 0)
		{
			objectNormal = nl;
			objectColor = hitColor;
			objectID = hitObjectID;
		}
		if (bounces == 1 && diffuseCount == 0 && !coatTypeIntersected)
		{
			objectNormal = nl;
		}


		// if we get here and sampleLight is still true, shadow ray failed to find a light source
		if (sampleLight) 
			break;
		
		
		    
                if (hitType == SPEC)  // Ideal SPECULAR reflection
                {
			mask *= hitColor;

			rayDirection = randomDirectionInSpecularLobe(reflect(rayDirection, nl), roughness);
			rayOrigin = x + nl * uEPS_intersect;
                        continue;
                }


                if (hitType == REFR)  // Ideal dielectric REFRACTION
		{
			pixelSharpness = diffuseCount == 0 ? -1.0 : pixelSharpness;

			nc = 1.0; // IOR of Air
			nt = 1.5; // IOR of common Glass
			Re = calcFresnelReflectance(rayDirection, n, nc, nt, ratioIoR);
			Tr = 1.0 - Re;
			P  = 0.25 + (0.5 * Re);
                	RP = Re / P;
                	TP = Tr / (1.0 - P);

			if (diffuseCount == 0 && rand() < P)
			{
				mask *= RP;
				rayDirection = randomDirectionInSpecularLobe(reflect(rayDirection, nl), roughness);
				rayOrigin = x + nl * uEPS_intersect;
                        	continue;
			}
			
			// transmit ray through surface
			mask *= TP;

			// is ray leaving a solid object from the inside? 
			// If so, attenuate ray color with object color by how far ray has travelled through the medium
			if (isRayExiting || (distance(n, nl) > 0.1))
			{
				isRayExiting = false;
				mask *= exp(log(hitColor) * thickness * t);
			}
			else 
				mask *= hitColor;
			
			tdir = refract(rayDirection, nl, ratioIoR);
			rayDirection = randomDirectionInSpecularLobe(tdir, roughness * roughness);
			rayOrigin = x - nl * uEPS_intersect;

			continue;
			
		} // end if (hitType == REFR)

		
		if (hitType == COAT)  // Diffuse object underneath with ClearCoat on top (like car, or shiny pool ball)
		{
			coatTypeIntersected = true;

			nc = 1.0; // IOR of Air
			nt = 1.5; // IOR of Clear Coat
			Re = calcFresnelReflectance(rayDirection, nl, nc, nt, ratioIoR);
			Tr = 1.0 - Re;
			P  = 0.25 + (0.5 * Re);
                	RP = Re / P;
                	TP = Tr / (1.0 - P);

			if (diffuseCount == 0 && rand() < P)
			{
				mask *= RP;
				rayDirection = randomDirectionInSpecularLobe(reflect(rayDirection, nl), roughness);
				rayOrigin = x + nl * uEPS_intersect;
                        	continue;
			}

			diffuseCount++;

			mask *= hitColor;
			mask *= TP;
			
			bounceIsSpecular = false;
 		
			
			if (diffuseCount == 1 && rand() < 0.5)
			{
				mask *= 2.0;	
				// choose random Diffuse sample vector
				rayDirection = randomCosWeightedDirectionInHemisphere(nl);
				rayOrigin = x + nl * uEPS_intersect;
				continue;
			} 
			
			randomLightChoice = rng(); // will be in the range 0.0-1.0
			// in Monte Carlo path tracing, we can only select 1 possible 'dirToLight' path by random choice
			// in other words, we can...

			// either sample the Sun
			if (randomLightChoice < 0.33)
			{
				dirToLight = randomDirectionInSpecularLobe(uSunDirectionVector, 0.03);
				weight = max(0.0, dot(dirToLight, nl)) * 0.000015; // down-weight directSunLight contribution
			}
			
			// or sample the large quad area light over the model
			else if (randomLightChoice < 0.66)
			{
				dirToLight = sampleQuadLight(x, nl, quads[0], weight);
			}
			
			// or sample the small sphere area light
			else if (randomLightChoice < 1.0)
			{
				dirToLight = sampleSphereLight(x, nl, spheres[0], weight);
			}
			
			// since we only selected 1 light source by random choice, but there are 3 light sources (much brighter)..
			// we must up-weight the contribution of the light that we did end up picking 
			weight *= 3.0; // 3.0 = number of light source choices (Sun, Quad, Sphere)
			// the following line also upweights because there was 0.5 chance that we reflect off of clearCoat, or sample diffuse surface beneath the clearCoat		 
			mask *= diffuseCount == 1 ? 2.0 : 1.0; // multiply by number of choices: 2.0 (either spec reflection, or diffuse)
			mask *= weight;

			rayDirection = dirToLight;
			rayOrigin = x + nl * uEPS_intersect;

			sampleLight = true;
			continue;
			
		} //end if (hitType == COAT)
		
	} // end for (int bounces = 0; bounces < 5; bounces++)
	

	return max(vec3(0), accumCol);

} // end vec3 CalculateRadiance( out vec3 objectNormal, out vec3 objectColor, out float objectID, out float pixelSharpness )


//-----------------------------------------------------------------------
void SetupScene(void)
//-----------------------------------------------------------------------
{
	vec3 z  = vec3(0);          
	vec3 L1 = vec3(5.0, 1.0, 1.0) * 25.0;// Bright Yellow light 
	vec3 L2 = vec3(6.0, 1.0, 0.0) * 1.0;// Yellow light 
	// spheres[0] = Sphere(  4000.0, vec3(0, -4000, 0),  z, vec3(0.4,0.4,0.4), CHECK);//Checkered Floor
	// spheres[1] = Sphere(     6.0, vec3(55, 36, -45),  z,         vec3(0.9),  SPEC);//small mirror ball
	// spheres[2] = Sphere(     6.0, vec3(55, 24, -45),  z, vec3(0.5,1.0,1.0),  REFR);//small glass ball
	// spheres[3] = Sphere(     6.0, vec3(60, 24, -30),  z,         vec3(1.0),  COAT);//small plastic ball
    	// spheres[0]  = Sphere(150.0, vec3(-400, 900, 200), L1, z, 0.0, SPEC);//spherical white Light1 
    	spheres[0] = Sphere( 3.2, vec3(0, 36, -45), L1, z, LIGHT);//small sphere light
	quads[0] = Quad( vec3(0,-1, 0), vec3(-10, 40,-32), vec3(3, 40,-32), vec3(3, 40,-27), vec3(-10, 40,-27), L2, z, LIGHT);// Area Light Rectangle above model
	// boxes[0] = Box( vec3(-20.0,11.0,-110.0), vec3(70.0,18.0,-20.0), z, vec3(0.2,0.9,0.7), REFR);//Glass Box
	// boxes[1] = Box( vec3(-14.0,13.0,-104.0), vec3(64.0,16.0,-26.0), z, vec3(0),           DIFF);//Inner Box
}


#include <pathtracing_main>