// We need only need one bounce given that we want to see the objects and then direct lighting is not done using raytracing
#pragma max_recursion_depth 31

#define HAS_LIGHTLOOP

// Include and define the shader pass
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/ShaderPass/ShaderPass.cs.hlsl"
#define SHADERPASS SHADERPASS_RAYTRACING

// HDRP include
#define SHADER_TARGET 50
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Color.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/Lighting.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/Material.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/LightLoop/LightLoopDef.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/NormalBuffer.hlsl"

#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/BSDF.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/PreIntegratedFGD/PreIntegratedFGD.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/CommonLighting.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/ImageBasedLighting.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/LightLoop/LightLoopDef.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/LightEvaluation.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/MaterialEvaluation.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariablesGlobal.hlsl"

// Raytracing includes
#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/ShaderVariablesRaytracingLightLoop.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/RayTracingLightCluster.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/RayTracingCommon.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/Lighting/ScreenSpaceLighting/ScreenSpaceLighting.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/Common/AtmosphericScatteringRayTracing.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Debug/RayCountManager.cs.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/RayTracingFallbackHierarchy.cs.hlsl"

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

// Flag value that defines if a given pixel recieves reflections or not
TEXTURE2D_X_UINT2(_StencilTexture);
int _SsrStencilBit;

// Output structure of the reflection raytrace shader
RW_TEXTURE2D_X(float4, _RayTracingLightingTextureRW);
RW_TEXTURE2D_X(float, _RayTracingDistanceTextureRW);

[shader("miss")]
void MissShaderReflections(inout RayIntersection rayIntersection : SV_RayPayload)
{
    float3 rayOrigin = WorldRayOrigin();
    float3 rayDirection = WorldRayDirection();

    float weight = 0.0f;

    if (RAYTRACINGFALLBACKHIERACHY_REFLECTION_PROBES & _RayTracingRayMissFallbackHierarchy)
        rayIntersection.color = RayTraceReflectionProbes(rayOrigin, rayDirection, weight);

    if((RAYTRACINGFALLBACKHIERACHY_SKY & _RayTracingRayMissFallbackHierarchy) && weight < 1.0)
    {
        rayIntersection.color += SAMPLE_TEXTURECUBE_ARRAY_LOD(_SkyTexture, s_trilinear_clamp_sampler, rayDirection, 0.0, 0).xyz * (1.0 - weight);
        weight = 1.0f;
    }

    if (weight > 0.0)
        ApplyFogAttenuation(rayOrigin, rayDirection, rayIntersection.color);

    rayIntersection.t = _RaytracingRayMaxLength;
}

// Function that evaluates the normal data for a given pixel
NormalData EvaluateNormalData(uint2 sourceCoord)
{
    // Decode the world space normal
    NormalData normalData;
    DecodeFromNormalBuffer(sourceCoord, normalData);

    // Override the roughness by the clearcoat value of this is a clear coat
    float4 coatMask = LOAD_TEXTURE2D_X(_SsrClearCoatMaskTexture, sourceCoord);
    normalData.perceptualRoughness = HasClearCoatMask(coatMask) ? CLEAR_COAT_SSR_PERCEPTUAL_ROUGHNESS : normalData.perceptualRoughness;
    return normalData;
}

[shader("raygeneration")]
void RayGenIntegration()
{
    // Grab the dimensions of the current dispatch
    uint3 LaunchIndex = DispatchRaysIndex();
    uint3 LaunchDim = DispatchRaysDimensions();

    UNITY_XR_ASSIGN_VIEW_INDEX(LaunchIndex.z);

    // Compute the pixel coordinate to evaluate
    uint2 currentCoord = uint2(LaunchIndex.x, LaunchDim.y - LaunchIndex.y - 1);

    // Clear the output color texture
    _RayTracingLightingTextureRW[COORD_TEXTURE2D_X(currentCoord)] = float4(0.0, 0.0, 0.0, 0.0);
    _RayTracingDistanceTextureRW[COORD_TEXTURE2D_X(currentCoord)] = 0.0;

    // Read the depth value
    float depthValue = LOAD_TEXTURE2D_X(_DepthTexture, currentCoord).r;

    // This point is part of the background, we don't really care
    if (depthValue == UNITY_RAW_FAR_CLIP_VALUE)
        return;

    // Does this pixel have SSS?
    uint stencilValue = GetStencilValue(LOAD_TEXTURE2D_X(_StencilTexture, currentCoord));
    if ((stencilValue & _SsrStencilBit) == 0)
        return;

    // Convert this to a world space position
    PositionInputs posInput = GetPositionInput(currentCoord, 1.0/LaunchDim.xy, depthValue, UNITY_MATRIX_I_VP, GetWorldToViewMatrix(), 0);
    float distanceToCamera = length(posInput.positionWS);
    // Compute the incident vector on the surfaces
    const float3 viewWS = GetWorldSpaceNormalizeViewDir(posInput.positionWS);

    // Decode the world space normal
    NormalData normalData = EvaluateNormalData(currentCoord);

    // Create the local ortho basis
    float3x3 localToWorld = GetLocalFrame(normalData.normalWS);

    // If this value is beyond the smothness that we allow, no need to compute it
    float perceptualSmoothness = PerceptualRoughnessToPerceptualSmoothness(normalData.perceptualRoughness);
    if (_RaytracingReflectionMinSmoothness > perceptualSmoothness)
        return;

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

    // Compute the actual roughness
    float roughness = PerceptualRoughnessToRoughness(normalData.perceptualRoughness);

    // If we only have one bounce, we don't need more than one ray to evaluate the exact signal. However, if we are going for multiple bounces, we may need more, so we cannot clamp the sample count to 1.
    int realSampleCount = _RaytracingMaxRecursion == 1 ? (normalData.perceptualRoughness < 0.01 ? 1 : _RaytracingNumSamples) : _RaytracingNumSamples;

    // Variable that accumulate the radiance
    float3 finalColor = float3(0.0, 0.0, 0.0);
    float averageDistance = 0.0;

    // Loop through the samples and add their contribution
    for (int sampleIndex = 0; sampleIndex < realSampleCount; ++sampleIndex)
    {
        // Compute the current sample index
        int globalSampleIndex = _RayTracingReflectionFrameIndex * realSampleCount + sampleIndex;

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

        // Importance sample the direction using GGX
        float3 sampleDir = float3(0.0, 0.0, 0.0);
        float NdotL, NdotH, VdotH;
        SampleGGXDir(theSample, viewWS, localToWorld, roughness, sampleDir, NdotL, NdotH, VdotH);

        // If the sample is under the surface
        if (dot(sampleDir, normalData.normalWS) <= 0.0)
            continue;

        // Make sure the new ray is taken into account in the ray counters
        if (_RayCountEnabled > 0)
        {
            uint3 counterIdx = uint3(currentCoord, INDEX_TEXTURE2D_ARRAY_X(RAYCOUNTVALUES_REFLECTION_FORWARD));
            _RayCountTexture[counterIdx] = _RayCountTexture[counterIdx] + 1;
        }

        // Create the ray descriptor for this pixel
        RayDesc rayDescriptor;
        rayDescriptor.Origin = posInput.positionWS + normalData.normalWS * rayBias;
        rayDescriptor.Direction = sampleDir;
        rayDescriptor.TMin = 0.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 rayIntersection;
        rayIntersection.color = float3(0.0, 0.0, 0.0);
        rayIntersection.t = -1.0;
        rayIntersection.remainingDepth = 1;
        rayIntersection.sampleIndex = globalSampleIndex;
        rayIntersection.pixelCoord = currentCoord;

        // In order to achieve filtering for the textures, we need to compute the spread angle of the pixel
        rayIntersection.cone.spreadAngle = _RaytracingPixelSpreadAngle + roughnessToSpreadAngle(roughness);
        rayIntersection.cone.width = distanceToCamera * _RaytracingPixelSpreadAngle;

        // Evaluate the ray intersection
        TraceRay(_RaytracingAccelerationStructure, RAY_FLAG_CULL_BACK_FACING_TRIANGLES, RAYTRACINGRENDERERFLAG_REFLECTION, 0, 1, 0, rayDescriptor, rayIntersection);

        // The clamping process is different for the sky and rest
        float3 sampleColor = 0.0;
        if (rayIntersection.t == _RaytracingRayMaxLength)
        {
            // Convert to HSV space
            sampleColor = RgbToHsv(rayIntersection.color * GetCurrentExposureMultiplier());

            // Expose and clamp the final color
            sampleColor.z = clamp(sampleColor.z, 0.0, _RaytracingIntensityClamp);

            // Convert back to HSV space
            sampleColor = HsvToRgb(sampleColor) * GetInverseCurrentExposureMultiplier();
        }
        else
        {
            // Expose and clamp the final color
            sampleColor = clamp(rayIntersection.color * GetCurrentExposureMultiplier(), 0.0, _RaytracingIntensityClamp) * GetInverseCurrentExposureMultiplier();
        }

        // Contribute to the pixel
        finalColor += sampleColor;
        averageDistance += rayIntersection.t;
    }

   	// Normalize the value
    if (realSampleCount != 0.0)
    {
        float normalizationFactor = 1.0 / realSampleCount;
        finalColor *= normalizationFactor;
        averageDistance *= normalizationFactor;
    }

    // We also need to compute the fade factor for this sample
    float weightValue = _RaytracingReflectionSmoothnessFadeStart == _RaytracingReflectionMinSmoothness ? 1.0 : saturate((perceptualSmoothness - _RaytracingReflectionMinSmoothness) / (_RaytracingReflectionSmoothnessFadeStart -_RaytracingReflectionMinSmoothness));

    // We store the sampled color and the weight that shall be used for it (1.0)
    _RayTracingLightingTextureRW[COORD_TEXTURE2D_X(currentCoord)] = float4(finalColor, weightValue);
    _RayTracingDistanceTextureRW[COORD_TEXTURE2D_X(currentCoord)] = averageDistance;
}

[shader("raygeneration")]
void RayGenIntegrationTransparent()
{
    // Grab the dimensions of the current dispatch
    uint3 LaunchIndex = DispatchRaysIndex();
    uint3 LaunchDim = DispatchRaysDimensions();

    UNITY_XR_ASSIGN_VIEW_INDEX(LaunchIndex.z);

    // Compute the pixel coordinate to evaluate
    uint2 currentCoord = uint2(LaunchIndex.x, LaunchDim.y - LaunchIndex.y - 1);

    // Clear the output color texture
    _RayTracingLightingTextureRW[COORD_TEXTURE2D_X(currentCoord)] = float4(0.0, 0.0, 0.0, 0.0);

    // Read the depth value
    float depthValue  = LOAD_TEXTURE2D_X(_DepthTexture, currentCoord).r;

    // This point is part of the background, we don't really care
    if (depthValue == UNITY_RAW_FAR_CLIP_VALUE)
        return;

    // Does this pixel have SSS?
    uint stencilValue = GetStencilValue(LOAD_TEXTURE2D_X(_StencilTexture, currentCoord));
    if ((stencilValue & _SsrStencilBit) == 0)
        return;

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

    float3 positionWS = posInput.positionWS;

    // Compute the incident vector on the surfaces
    const float3 viewWS = GetWorldSpaceNormalizeViewDir(posInput.positionWS);

    // Decode the world space normal
    NormalData normalData;
    DecodeFromNormalBuffer(currentCoord, normalData);

    // If this value is beyond the smothness that we allow, no need to compute it
    float perceptualSmoothness = PerceptualRoughnessToPerceptualSmoothness(normalData.perceptualRoughness);
    if (_RaytracingReflectionMinSmoothness > perceptualSmoothness)
        return;

    // Compute the reflected direction
    float3 reflectionDir = reflect(-viewWS, normalData.normalWS);

    // Make sure the new ray is taken into account in the ray counters
    if (_RayCountEnabled > 0)
    {
        uint3 counterIdx = uint3(currentCoord, INDEX_TEXTURE2D_ARRAY_X(RAYCOUNTVALUES_REFLECTION_FORWARD));
        _RayCountTexture[counterIdx] = _RayCountTexture[counterIdx] + 1;
    }

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

    // Create the ray descriptor for this pixel
    RayDesc rayDescriptor;
    rayDescriptor.Origin = positionWS + normalData.normalWS * rayBias;
    rayDescriptor.Direction = reflectionDir;
    rayDescriptor.TMin = 0.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 rayIntersection;
    rayIntersection.color = float3(0.0, 0.0, 0.0);
    rayIntersection.t = -1.0;
    rayIntersection.remainingDepth = 1;
    rayIntersection.sampleIndex = 0;
    rayIntersection.pixelCoord = currentCoord;

    // In order to achieve filtering for the textures, we need to compute the spread angle of the pixel
    rayIntersection.cone.spreadAngle = _RaytracingPixelSpreadAngle;
    rayIntersection.cone.width = distanceToCamera * _RaytracingPixelSpreadAngle;

    // Evaluate the ray intersection
    TraceRay(_RaytracingAccelerationStructure, RAY_FLAG_CULL_BACK_FACING_TRIANGLES, RAYTRACINGRENDERERFLAG_REFLECTION, 0, 1, 0, rayDescriptor, rayIntersection);

    // We store the sampled color and the weight that shall be used for it (1.0)
    _RayTracingLightingTextureRW[COORD_TEXTURE2D_X(currentCoord)] = float4(rayIntersection.color, 1.0);
    _RayTracingDistanceTextureRW[COORD_TEXTURE2D_X(currentCoord)] = rayIntersection.t;
}