#pragma kernel RaytracingReflectionsFullRes
#pragma kernel RaytracingReflectionsHalfRes

#pragma kernel RaytracingReflectionsTransparentFullRes
#pragma kernel RaytracingReflectionsTransparentHalfRes

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

#pragma only_renderers d3d11 xboxseries ps5

// #pragma enable_d3d11_debug_symbols

// HDRP generic includes
#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/Material/Material.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Material/NormalBuffer.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/ScreenSpaceLighting/ScreenSpaceLighting.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/RaytracingSampling.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/RenderPipeline/Raytracing/Shaders/RayTracingCommon.hlsl"

// Tile size of this compute
#define RAYTRACING_REFLECTIONS_TILE_SIZE 8

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

// Output data
RW_TEXTURE2D_X(float4, _RaytracingDirectionBuffer);

// Structure that holds everything that defines a GGX sample
struct GGXSample
{
    float3 direction;
    float pdf;
};

GGXSample GenerateGGXSampleDirection(uint2 currentCoord, float3 normalWS, float3 viewWS, float roughness, int frameIndex)
{
    // Create the local ortho basis
    float3x3 localToWorld = GetLocalFrame(normalWS);

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

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

    // If this direction is under the surface, let's generate a new one that won't be.
    // We allow ourselves 8 as the total of number of tries.
    // TODO: use Eric's paper on visible normal distribution sampling
    initialFrameIndex += 8;
    for (int i = 1; i < 8; ++i)
    {
        if (dot(sampleDir, normalWS) >= 0.00f)
            break;

        theSample.x = GetBNDSequenceSample(currentCoord, initialFrameIndex + i, 0);
        theSample.y = GetBNDSequenceSample(currentCoord, initialFrameIndex + i, 1);
        SampleGGXDir(theSample, viewWS, localToWorld, roughness, sampleDir, NdotL, NdotH, VdotH);
    }

    // Build the sample and return it
    GGXSample ggxSample;
    ggxSample.direction = sampleDir;
    // Given that GGX is invalid for a pure smooth material, we handle the case this by stating that the pdf == 1.0
    ggxSample.pdf = roughness > 0.001 ? D_GGX(NdotH, roughness) * NdotH / (4.0 * VdotH) : 1.0;
    return ggxSample;
}

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

[numthreads(RAYTRACING_REFLECTIONS_TILE_SIZE, RAYTRACING_REFLECTIONS_TILE_SIZE, 1)]
void RaytracingReflectionsHalfRes(uint3 dispatchThreadId : SV_DispatchThreadID, uint2 groupThreadId : SV_GroupThreadID, uint2 groupId : SV_GroupID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);

    // Compute the pixel position to process
    uint2 halfResCoord = groupId * RAYTRACING_REFLECTIONS_TILE_SIZE + groupThreadId;

    // Pixel coordinate in full res that we will be using to pick depth, normal and smoothness values.
    uint2 sourceCoord = ComputeSourceCoordinates(halfResCoord, _RayTracingCheckerIndex);

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

    // We initialize the output texture in case we early return. The last channel is set to -1.0 to mark the invalidity of the direction.
    _RaytracingDirectionBuffer[COORD_TEXTURE2D_X(sourceCoord)] = float4(0.0, 0.0, 0.0, -1.0);

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

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

    // Convert this to a world space position
    PositionInputs posInput = GetPositionInput(sourceCoord, _ScreenSize.zw, depthValue, UNITY_MATRIX_I_VP, GetWorldToViewMatrix(), 0);

    // Compute the view vector
    const float3 viewWS = GetWorldSpaceNormalizeViewDir(posInput.positionWS);

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

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

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

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

    // Generate a GGX direction
    GGXSample ggxSample = GenerateGGXSampleDirection(halfResCoord, normalData.normalWS, viewWS, roughness, _RayTracingReflectionFrameIndex);

    // If we failed, return.
    if (dot(ggxSample.direction, normalData.normalWS) <= 0.00f)
        return;

    // Store the generated direction and inverted PDF
    _RaytracingDirectionBuffer[COORD_TEXTURE2D_X(sourceCoord)] = float4(ggxSample.direction, 1.0 / ggxSample.pdf);
}

[numthreads(RAYTRACING_REFLECTIONS_TILE_SIZE, RAYTRACING_REFLECTIONS_TILE_SIZE, 1)]
void RaytracingReflectionsTransparentHalfRes(uint3 dispatchThreadId : SV_DispatchThreadID, uint2 groupThreadId : SV_GroupThreadID, uint2 groupId : SV_GroupID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);

    // Compute the pixel position to process
    uint2 halfResCoord = groupId * RAYTRACING_REFLECTIONS_TILE_SIZE + groupThreadId;

    // Pixel coordinate in full res that we will be using to pick depth and normal
    uint2 sourceCoord = ComputeSourceCoordinates(halfResCoord, _RayTracingCheckerIndex);

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

    // We initialize the output texture in case we early return. The last channel is set to -1.0 to mark the invalidity of the direction.
    _RaytracingDirectionBuffer[COORD_TEXTURE2D_X(sourceCoord)] = float4(0.0, 0.0, 0.0, -1.0);

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

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

    // Convert this to a world space position
    PositionInputs posInput = GetPositionInput(sourceCoord, _ScreenSize.zw, depthValue, UNITY_MATRIX_I_VP, GetWorldToViewMatrix(), 0);

    // Compute the view vector
    const float3 viewWS = GetWorldSpaceNormalizeViewDir(posInput.positionWS);

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

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

    // Store the reflected direction and 1.0 in the w channel.
    _RaytracingDirectionBuffer[COORD_TEXTURE2D_X(sourceCoord)] = float4(reflect(-viewWS, normalData.normalWS), 1.0);
}

[numthreads(RAYTRACING_REFLECTIONS_TILE_SIZE, RAYTRACING_REFLECTIONS_TILE_SIZE, 1)]
void RaytracingReflectionsFullRes(uint3 dispatchThreadId : SV_DispatchThreadID, uint2 groupThreadId : SV_GroupThreadID, uint2 groupId : SV_GroupID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);

    // Compute the pixel position to process
    uint2 currentCoord = groupId * RAYTRACING_REFLECTIONS_TILE_SIZE + groupThreadId;

    // Clear the output color texture
    _RaytracingDirectionBuffer[COORD_TEXTURE2D_X(currentCoord)] = float4(0.0, 0.0, 0.0, -1.0f);

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

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

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

    // Convert this to a world space position
    PositionInputs posInput = GetPositionInput(currentCoord, _ScreenSize.zw, depthValue, UNITY_MATRIX_I_VP, GetWorldToViewMatrix(), 0);

    // Compute the view vector
    const float3 viewWS = GetWorldSpaceNormalizeViewDir(posInput.positionWS);

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

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

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

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

    // Generate a GGX direction
    GGXSample ggxSample = GenerateGGXSampleDirection(currentCoord, normalData.normalWS, viewWS, roughness, _RayTracingReflectionFrameIndex);

    // If we failed, return.
    if (dot(ggxSample.direction, normalData.normalWS) <= 0.00f)
        return;

    // Store the generated direction and inverted PDF
    _RaytracingDirectionBuffer[COORD_TEXTURE2D_X(currentCoord)] = float4(ggxSample.direction, 1.0 / ggxSample.pdf);
}

[numthreads(RAYTRACING_REFLECTIONS_TILE_SIZE, RAYTRACING_REFLECTIONS_TILE_SIZE, 1)]
void RaytracingReflectionsTransparentFullRes(uint3 dispatchThreadId : SV_DispatchThreadID, uint2 groupThreadId : SV_GroupThreadID, uint2 groupId : SV_GroupID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(dispatchThreadId.z);

    // Compute the pixel position to process
    uint2 currentCoord = groupId * RAYTRACING_REFLECTIONS_TILE_SIZE + groupThreadId;

    // We initialize the output texture in case we early return. The last channel is set to -1.0 to mark the invalidity of the direction.
    _RaytracingDirectionBuffer[COORD_TEXTURE2D_X(currentCoord)] = float4(0.0, 0.0, 0.0, -1.0f);

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

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

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

    // Convert this to a world space position
    PositionInputs posInput = GetPositionInput(currentCoord, _ScreenSize.zw, depthValue, UNITY_MATRIX_I_VP, GetWorldToViewMatrix(), 0);

    // Compute the view vector
    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
    if (_RaytracingReflectionMinSmoothness > PerceptualRoughnessToPerceptualSmoothness(normalData.perceptualRoughness))
        return;

    // Write the output ray data
    _RaytracingDirectionBuffer[COORD_TEXTURE2D_X(currentCoord)] = float4(reflect(-viewWS, normalData.normalWS), 1.0);
}