// 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/HDStencilUsage.cs.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/RenderPipeline/Raytracing/Shaders/Common/RayTracingHelpers.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);
TEXTURE2D_X_UINT2(_StencilTexture);

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

[shader("miss")]
void MissShaderAmbientOcclusion(inout RayIntersectionVisibility 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;
    uint stencilValue = GetStencilValue(LOAD_TEXTURE2D_X(_StencilTexture, currentPixelCoord));
    // This point is part of the background or is unlit, we don't really care
    if (depthValue == UNITY_RAW_FAR_CLIP_VALUE || (stencilValue & STENCILUSAGE_IS_UNLIT) != 0)
		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;
    }

    // Evaluate the ray bias
    float rayBias = EvaluateRayTracingBias(posInput.positionWS);

	// Variable that accumulate the radiance
	float finalColor = 0.0;
	float velocity = 0.0;
	RayDesc rayDescriptor;
	RayIntersectionVisibility rayIntersection;

	// 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
		rayDescriptor.Origin = posInput.positionWS + normalData.normalWS * rayBias;
		rayDescriptor.Direction = sampleDir;
		rayDescriptor.TMin = 0;
		rayDescriptor.TMax = _RaytracingRayMaxLength;

		// Adjust world-space position to match the RAS setup with XR single-pass and camera relative
		ApplyCameraRelativeXR(rayDescriptor.Origin);

		// Create and init the RayIntersection structure for this
		rayIntersection.color = float3(0.0, 0.0, 0.0);
		rayIntersection.t = 0.0;
        rayIntersection.pixelCoord = posInput.positionSS;
        rayIntersection.velocity = 0.0;

		// 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;
}