#version 460

#extension GL_EXT_ray_tracing : enable
#extension GL_ARB_shader_clock : enable

#pragma shader_stage(raygen)
layout (constant_id = 0) const uint PASSES = 2;
layout (constant_id = 1) const uint TEXTURES = 512;
layout (constant_id = 2) const uint CUBEMAPS = 8;
layout (constant_id = 3) const uint CASCADES = 1;

// shader type settings
#define RT 1
#define DEFERRED_SAMPLING 1
#define BUFFER_REFERENCE 1
#define UINT64_ENABLED 1
#define PBR 1
// shader function settings
#define FOG 1

//force it off
#define BARYCENTRIC 0
#define BARYCENTRIC_CALCULATE 0

#include "../common/macros.h"
#include "../common/structs.h"

layout( push_constant ) uniform PushBlock {
	uint pass;
	uint draw;
} PushConstant;

layout (binding = 0) uniform accelerationStructureEXT tlas;

layout (binding = 1, rgba32f) uniform volatile coherent image2D outImage;

layout (binding = 2) uniform UBO {
	EyeMatrices eyes[2];

	Settings settings;
} ubo;

layout (std140, binding = 3) readonly buffer Instances {
	Instance instances[];
};
layout (std140, binding = 4) readonly buffer InstanceAddresseses {
	InstanceAddresses instanceAddresses[];
};
layout (std140, binding = 5) readonly buffer Materials {
	Material materials[];
};
layout (std140, binding = 6) readonly buffer Textures {
	Texture textures[];
};
layout (std140, binding = 7) readonly buffer Lights {
	Light lights[];
};

layout (binding = 8) uniform sampler2D samplerTextures[TEXTURES];
layout (binding = 9) uniform samplerCube samplerCubemaps[CUBEMAPS];
layout (binding = 10) uniform sampler3D samplerNoise;
#if VXGI
	layout (binding = 11) uniform usampler3D voxelId[CASCADES];
	layout (binding = 12) uniform sampler3D voxelNormal[CASCADES];
	layout (binding = 13) uniform sampler3D voxelRadiance[CASCADES];
#endif

layout (location = 0) rayPayloadEXT RayTracePayload payload;

//layout(buffer_reference, scalar) buffer Vertices { Vertex v[]; };
layout(buffer_reference, scalar) buffer Indices { uint i[]; };
//layout(buffer_reference, scalar) buffer Indices { uvec3 i[]; };
layout(buffer_reference, scalar) buffer Indirects { DrawCommand dc[]; };

//layout(buffer_reference, scalar) buffer VPos { vec3 v[]; };
layout(buffer_reference, scalar) buffer VPos { float v[]; };
layout(buffer_reference, scalar) buffer VUv { vec2 v[]; };
//layout(buffer_reference, scalar) buffer VUv { float v[]; };
layout(buffer_reference, scalar) buffer VColor { uint v[]; };
layout(buffer_reference, scalar) buffer VSt { vec2 v[]; };
//layout(buffer_reference, scalar) buffer VSt { float v[]; };
//layout(buffer_reference, scalar) buffer VNormal { vec3 v[]; };
layout(buffer_reference, scalar) buffer VNormal { float v[]; };
//layout(buffer_reference, scalar) buffer VTangent { vec3 v[]; };
layout(buffer_reference, scalar) buffer VTangent { float v[]; };
layout(buffer_reference, scalar) buffer VID { uint v[]; };

#include "../common/functions.h"
#include "../common/light.h"
#include "../common/fog.h"
#if VXGI
	#include "../common/vxgi.h"
#endif

void trace( Ray ray, float tMin, float tMax ) {
	uint rayFlags = gl_RayFlagsOpaqueEXT;
	uint cullMask = 0xFF;
	
	payload.hit = false;
	surface.position.eye.z = tMax;
	traceRayEXT(tlas, rayFlags, cullMask, 0, 0, 0, ray.origin, tMin, ray.direction, tMax, 0);
}
void trace( Ray ray, float tMin ) {
	uint rayFlags = gl_RayFlagsOpaqueEXT;
	uint cullMask = 0xFF;

	float tMax = ubo.settings.rt.defaultRayBounds.y;
	
	payload.hit = false;
	surface.position.eye.z = tMax;
	traceRayEXT(tlas, rayFlags, cullMask, 0, 0, 0, ray.origin, tMin, ray.direction, tMax, 0);
}

void trace( Ray ray ) {
	trace( ray, ubo.settings.rt.defaultRayBounds.x, ubo.settings.rt.defaultRayBounds.y );
}

float shadowFactor( const Light light, float def ) {
	Ray ray;
	ray.origin = surface.position.world;
	ray.direction = light.position - ray.origin;

	float tMin = ubo.settings.rt.defaultRayBounds.x;
	float tMax = length(ray.direction) - 0.0001;
	
	ray.direction = normalize(ray.direction);

	uint rayFlags = gl_RayFlagsOpaqueEXT | gl_RayFlagsTerminateOnFirstHitEXT | gl_RayFlagsSkipClosestHitShaderEXT;
	uint cullMask = 0xFF;

	payload.hit = true;
	traceRayEXT(tlas, rayFlags, cullMask, 0, 0, 0, ray.origin, tMin, ray.direction, tMax, 0);

	return payload.hit ? 0.0 : 1.0;
}

void directLighting() {
#if VXGI
	indirectLighting();
#endif

	surface.light.rgb += surface.material.albedo.rgb * ubo.settings.lighting.ambient.rgb * surface.material.occlusion; // add ambient lighting
	surface.light.rgb += surface.material.indirect.rgb; // add indirect lighting
#if PBR
	pbr();
#elif LAMBERT
	lambert();
#elif PHONG
	phong();
#endif
	surface.fragment.rgb += surface.light.rgb;
	surface.fragment.a = surface.material.albedo.a;
}

vec4 traceStep( Ray ray ) {
	Ray fogRay = ray;
	float eyeDepth = 0;
	vec4 outFrag = vec4(0);

	// initial condition
	{
		trace( ray );

		if ( payload.hit ) {
			populateSurface( payload );
			directLighting();
		} else if (  0 <= ubo.settings.lighting.indexSkybox && ubo.settings.lighting.indexSkybox < CUBEMAPS ) {
			surface.fragment = texture( samplerCubemaps[ubo.settings.lighting.indexSkybox], ray.direction );
			surface.fragment.a = 4096;
			surface.position.eye.z /= 8;
		} else {
			surface.fragment = vec4(ubo.settings.lighting.ambient.rgb, 0.5);
		}
	#if FOG
		fog( ray, surface.fragment.rgb, surface.fragment.a );
	#endif
		outFrag = surface.fragment;
		eyeDepth = surface.position.eye.z;
	}


	// "transparency"
	if ( payload.hit && surface.material.albedo.a < 0.999 ) {
		const vec4 TRANSPARENCY_COLOR = vec4(1.0 - surface.material.albedo.a);
		
		if ( surface.material.albedo.a < 0.001 ) outFrag = vec4(0);

		RayTracePayload surfacePayload = payload;
		Ray transparency;
		transparency.direction = ray.direction;
		transparency.origin = surface.position.world;
		fogRay = transparency;

		trace( transparency, ubo.settings.rt.alphaTestOffset );
		if ( payload.hit ) {
			populateSurface( payload );
			directLighting();
		} else if (  0 <= ubo.settings.lighting.indexSkybox && ubo.settings.lighting.indexSkybox < CUBEMAPS ) {
			surface.fragment = texture( samplerCubemaps[ubo.settings.lighting.indexSkybox], ray.direction );
			surface.fragment.a = 4096;
			surface.position.eye.z /= 8;
		}
	#if FOG
		fog( transparency, surface.fragment.rgb, surface.fragment.a );
	#endif
		outFrag += TRANSPARENCY_COLOR * surface.fragment;
		eyeDepth = surface.position.eye.z;

		payload = surfacePayload;
		populateSurface( payload );
	}
#if FOG
	{
	//	surface.position.eye.z = eyeDepth;
	//	fog( fogRay, outFrag.rgb, outFrag.a );
	//	fog( ray, surface.fragment.rgb, surface.fragment.a );
	}
#endif

	// reflection
	if ( payload.hit ) {
		const float REFLECTIVITY = 1.0 - surface.material.roughness;
		const vec4 REFLECTED_ALBEDO = surface.material.albedo * REFLECTIVITY;

		if ( REFLECTIVITY > 0.001 ) {
			RayTracePayload surfacePayload = payload;

			Ray reflection;
			reflection.origin = surface.position.world;
			reflection.direction = reflect( ray.direction, surface.normal.world );

			trace( reflection );

			if ( payload.hit ) {
				populateSurface( payload );
				directLighting();
			} else if (  0 <= ubo.settings.lighting.indexSkybox && ubo.settings.lighting.indexSkybox < CUBEMAPS ) {
				surface.fragment = texture( samplerCubemaps[ubo.settings.lighting.indexSkybox], reflection.direction );
				surface.fragment.a = 4096;
			}
		#if FOG
			fog( reflection, surface.fragment.rgb, surface.fragment.a );
		#endif
			outFrag += REFLECTED_ALBEDO * surface.fragment;

			payload = surfacePayload;
			populateSurface( payload );
		}
	}

	return outFrag;
}

void main()  {
//	if ( ubo.settings.mode.frameNumber > 16 ) return;	
//	prngSeed = tea(gl_LaunchIDEXT.y * gl_LaunchSizeEXT.x + gl_LaunchIDEXT.x, ubo.settings.mode.frameNumber);
	prngSeed = tea(gl_LaunchIDEXT.y * gl_LaunchSizeEXT.x + gl_LaunchIDEXT.x, int(clockARB()));
	surface.pass = PushConstant.pass;
	surface.subID = 0;
	surface.fragment = vec4(0);
	surface.light = vec4(0);
	surface.motion = vec2(0);
	surface.material.indirect = vec4(0);
	
	vec4 outFrag = vec4(0);

	const uint SAMPLES = min(ubo.settings.rt.samples, 4);
	const uint NUM_PATHS = min(ubo.settings.rt.paths, 8);
#if 1
	const uint FRAME_ACCUMULATION_VALUE = ubo.settings.rt.frameAccumulationMinimum > 0 ? min(ubo.settings.rt.frameAccumulationMinimum, ubo.settings.mode.frameNumber + 1) : ubo.settings.mode.frameNumber + 1;
#else
	const uint FRAME_ACCUMULATION_VALUE = min(32, ubo.settings.mode.frameNumber + 1);
#endif
	const float BLEND_FACTOR = 1.0f / float(FRAME_ACCUMULATION_VALUE);
	uint FRAME_NUMBER = ubo.settings.mode.frameNumber;

#if 0
	for ( uint samp = 0; samp < SAMPLES; ++samp, ++FRAME_NUMBER ) {
		{
			const vec2 center = ( FRAME_NUMBER > 0 ) ? vec2( rnd(), rnd() ) : vec2(0.5);
			const vec2 inUv = (vec2(gl_LaunchIDEXT.xy) + center) / vec2(gl_LaunchSizeEXT.xy);
		#if 1
			vec4 target	= ubo.eyes[surface.pass].iProjection * vec4(inUv.x * 2.0f - 1.0f, inUv.y * 2.0f - 1.0f, 1, 1);
			vec4 direction = ubo.eyes[surface.pass].iView * vec4(normalize(target.xyz), 0);

			surface.ray.direction = vec3(direction);
			surface.ray.origin = ubo.eyes[surface.pass].eyePos.xyz;
		#else
			const mat4 iProjectionView = inverse( ubo.eyes[surface.pass].projection * mat4(mat3(ubo.eyes[surface.pass].view)) );
			const vec4 near4 = iProjectionView * (vec4(2.0 * inUv - 1.0, -1.0, 1.0));
			const vec4 far4 = iProjectionView * (vec4(2.0 * inUv - 1.0, 1.0, 1.0));
			const vec3 near3 = near4.xyz / near4.w;
			const vec3 far3 = far4.xyz / far4.w;

			surface.ray.direction = normalize( far3 - near3 );
			surface.ray.origin = ubo.eyes[surface.pass].eyePos.xyz;
		#endif
		}

		{
			vec4 curValue = vec4(0);
			vec4 curWeight = vec4(1);
			for ( uint path = 0; path < NUM_PATHS; ++path ) {
				vec4 stepValue = traceStep( surface.ray );
				curValue += stepValue * curWeight;

				if ( !payload.hit ) break;

				surface.ray.origin = surface.position.world;
				surface.ray.direction = samplingHemisphere( prngSeed, surface.normal.world );
				curWeight *= surface.material.albedo * dot( surface.ray.direction, surface.normal.world );
				
				if ( length(curWeight) < 0.01 ) break;
			}
			outFrag += curValue;
		}
	}
	{
		outFrag /= SAMPLES;
	}
#elif 0
	{
		const vec2 center = ( FRAME_NUMBER > 0 ) ? vec2( rnd(), rnd() ) : vec2(0.5);
		const vec2 inUv = (vec2(gl_LaunchIDEXT.xy) + center) / vec2(gl_LaunchSizeEXT.xy);
	#if 0
		vec4 target	= ubo.eyes[surface.pass].iProjection * vec4(inUv.x * 2.0f - 1.0f, inUv.y * 2.0f - 1.0f, 1, 1);
		vec4 direction = ubo.eyes[surface.pass].iView * vec4(normalize(target.xyz), 0);

		surface.ray.direction = vec3(direction);
		surface.ray.origin = ubo.eyes[surface.pass].eyePos.xyz;
	#else
		const mat4 iProjectionView = inverse( ubo.eyes[surface.pass].projection * mat4(mat3(ubo.eyes[surface.pass].view)) );
		const vec4 near4 = iProjectionView * (vec4(2.0 * inUv - 1.0, -1.0, 1.0));
		const vec4 far4 = iProjectionView * (vec4(2.0 * inUv - 1.0, 1.0, 1.0));
		const vec3 near3 = near4.xyz / near4.w;
		const vec3 far3 = far4.xyz / far4.w;

		surface.ray.direction = normalize( far3 - near3 );
		surface.ray.origin = ubo.eyes[surface.pass].eyePos.xyz;
	#endif
	}

	{
		vec4 curValue = vec4(0);
		vec4 curWeight = vec4(1);
		for ( uint path = 0; path < NUM_PATHS; ++path ) {
			vec4 stepValue = traceStep( surface.ray );
			curValue += stepValue * curWeight;

			if ( !payload.hit ) break;

			surface.ray.origin = surface.position.world;
			surface.ray.direction = samplingHemisphere( prngSeed, surface.normal.world );
			curWeight *= surface.material.albedo * dot( surface.ray.direction, surface.normal.world );
			
			if ( length(curWeight) < 0.01 ) break;
		}
		outFrag += curValue;
	}
	{
		surface.fragment = outFrag;
	}
#else
	{
		const vec2 center = ( FRAME_NUMBER > 0 ) ? vec2( rnd(), rnd() ) : vec2(0.5);
		const vec2 inUv = (vec2(gl_LaunchIDEXT.xy) + center) / vec2(gl_LaunchSizeEXT.xy);
	#if 0
		vec4 target	= ubo.eyes[surface.pass].iProjection * vec4(inUv.x * 2.0f - 1.0f, inUv.y * 2.0f - 1.0f, 1, 1);
		vec4 direction = ubo.eyes[surface.pass].iView * vec4(normalize(target.xyz), 0);

		surface.ray.direction = vec3(direction);
		surface.ray.origin = ubo.eyes[surface.pass].eyePos.xyz;
	#else
		const mat4 iProjectionView = inverse( ubo.eyes[surface.pass].projection * mat4(mat3(ubo.eyes[surface.pass].view)) );
		const vec4 near4 = iProjectionView * (vec4(2.0 * inUv - 1.0, -1.0, 1.0));
		const vec4 far4 = iProjectionView * (vec4(2.0 * inUv - 1.0, 1.0, 1.0));
		const vec3 near3 = near4.xyz / near4.w;
		const vec3 far3 = far4.xyz / far4.w;

		surface.ray.direction = normalize( far3 - near3 );
		surface.ray.origin = ubo.eyes[surface.pass].eyePos.xyz;
	#endif
	}
	{
		surface.fragment = traceStep( surface.ray );
	}
#endif
	{
	#if BLOOM
		float brightness = dot(surface.fragment.rgb, vec3(0.2126, 0.7152, 0.0722));
		vec4 outFragBright = brightness > ubo.threshold ? vec4(surface.fragment.rgb, 1.0) : vec4(0, 0, 0, 1);
	//	imageStore(outImage, ivec2(gl_LaunchIDEXT.xy), outFragBright);
	#endif
	#if FOG
		fog( surface.ray, surface.fragment.rgb, surface.fragment.a );
	#endif
	}

	{
		outFrag = surface.fragment;
		outFrag.a = 1;
	}

	if ( ubo.settings.mode.frameNumber == 0 ) {
		imageStore(outImage, ivec2(gl_LaunchIDEXT.xy), outFrag);
	} else {
	//	if ( length(outFrag.rgb) < 0.01f ) return;
		vec4 blended = mix(imageLoad(outImage, ivec2(gl_LaunchIDEXT.xy)), outFrag, BLEND_FACTOR);

		imageStore(outImage, ivec2(gl_LaunchIDEXT.xy), blended);
	}
}