// We need only need 1 bounce for AO
#pragma max_recursion_depth 1

// HDRP include
#define SHADER_TARGET 50
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Macros.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Packing.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariablesFunctions.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Sampling/Sampling.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/NormalBuffer.hlsl"

#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/ShaderVariablesRaytracing.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/RaytracingIntersection.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/RaytracingSampling.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Debug/RayCountManager.cs.hlsl"

// The target acceleration structure that we will evaluate the reflexion in
TEXTURE2D_X(_DepthTexture);

// Output structure of the reflection raytrace shader
RW_TEXTURE2D_X(float, _AmbientOcclusionTextureRW);
RW_TEXTURE2D_X(float, _VelocityBuffer);

[shader("miss")]
void MissShaderAmbientOcclusion(inout RayIntersection rayIntersection : SV_RayPayload)
{
	rayIntersection.color += float3(1.0f, 1.0f, 1.0f);
}

[shader("raygeneration")]
void RayGenAmbientOcclusion()
{
	uint3 LaunchIndex = DispatchRaysIndex();
    uint2 LaunchDim = DispatchRaysDimensions().xy;

    UNITY_XR_ASSIGN_VIEW_INDEX(LaunchIndex.z);

	// Pixel coordinate of the current pixel
    uint2 currentPixelCoord = uint2(LaunchIndex.x, LaunchIndex.y);

    // Reset the value of this pixel
    _AmbientOcclusionTextureRW[COORD_TEXTURE2D_X(currentPixelCoord)] = 0.0f;
	
	// Read the depth value
	float depthValue  = LOAD_TEXTURE2D_X(_DepthTexture, currentPixelCoord).r;
	if (depthValue == UNITY_RAW_FAR_CLIP_VALUE)
		return;

	// Convert this to a world space position
	PositionInputs posInput = GetPositionInput(currentPixelCoord, 1.0/LaunchDim.xy, depthValue, UNITY_MATRIX_I_VP, GetWorldToViewMatrix(), 0);

	// Decode the world space normal
    NormalData normalData;	
    DecodeFromNormalBuffer(currentPixelCoord, normalData);
	
	// the number of samples based on the roughness
	int numSamples = _RaytracingNumSamples;

    // Count the number of rays that we will be traced
    if (_RayCountEnabled > 0)
    {
        uint3 counterIdx = uint3(currentPixelCoord, INDEX_TEXTURE2D_ARRAY_X(RAYCOUNTVALUES_AMBIENT_OCCLUSION));
        _RayCountTexture[counterIdx] = _RayCountTexture[counterIdx] + (uint)numSamples;
    }

	// Variable that accumulate the radiance
	float finalColor = 0.0;
	float velocity = 0.0f;
	// Let's loop through th e samples
	for (int i = 0; i < numSamples; ++i)
	{
        // Compute the current sample index
        int globalSampleIndex = _RaytracingFrameIndex * _RaytracingNumSamples + i;

    	// Generate the new sample (follwing values of the sequence)
        float2 noiseValue;
        noiseValue.x = GetBNDSequenceSample(currentPixelCoord, globalSampleIndex, 0);
        noiseValue.y = GetBNDSequenceSample(currentPixelCoord, globalSampleIndex, 1);

    	// Importance sample the direction
	    float3 sampleDir = SampleHemisphereCosine(noiseValue.x, noiseValue.y, normalData.normalWS);
		
		// Create the ray descriptor for this pixel
		RayDesc rayDescriptor;
		rayDescriptor.Origin = posInput.positionWS + normalData.normalWS * _RaytracingRayBias;
		rayDescriptor.Direction = sampleDir;
		rayDescriptor.TMin = 0;
		rayDescriptor.TMax = _RaytracingRayMaxLength;

		// Create and init the RayIntersection structure for this
		RayIntersection rayIntersection;
		rayIntersection.color = float3(0.0, 0.0, 0.0);
		rayIntersection.incidentDirection = rayDescriptor.Direction;
		rayIntersection.origin = rayDescriptor.Origin;
		rayIntersection.t = 0.0f;
        rayIntersection.pixelCoord = posInput.positionSS;

		// Evaluate the ray intersection
		TraceRay(_RaytracingAccelerationStructure, RAY_FLAG_CULL_BACK_FACING_TRIANGLES, RAYTRACINGRENDERERFLAG_AMBIENT_OCCLUSION, 0, 1, 0, rayDescriptor, rayIntersection);
			
		// Accumulate this value
		velocity = max(velocity, rayIntersection.velocity);
		finalColor += rayIntersection.color.x;
	}

	// Normalize the radiance
	finalColor /= (float)numSamples;

	// Alright we are done
    _AmbientOcclusionTextureRW[COORD_TEXTURE2D_X(currentPixelCoord)] = finalColor;
    _VelocityBuffer[COORD_TEXTURE2D_X(currentPixelCoord)] = velocity;
}

// Fallback default any hit shader for this raytrace shader
[shader("anyhit")]
void AnyHitMain(inout RayIntersection rayIntersection : SV_RayPayload, AttributeData attributeData : SV_IntersectionAttributes)
{
    rayIntersection.color = float3(0.0, 0.0, 0.0);
    AcceptHitAndEndSearch();
}