#pragma kernel BilateralUpSampleColorHalf
#pragma kernel BilateralUpSampleColor

//#pragma enable_d3d11_debug_symbols

#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch

#define BILATERAL_UPSAMPLE_TILE_SIZE 8

#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/RenderPipeline/Raytracing/Shaders/RayTracingCommon.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/ScreenSpaceLighting/BilateralUpsample.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Lighting/ScreenSpaceLighting/BilateralUpsampleDef.cs.hlsl"

// Mip chain depth buffer
TEXTURE2D_X(_DepthTexture);
// The half resolution texture that needs to be upscaled
TEXTURE2D_X(_LowResolutionTexture);

// LDS that store the half resolution data
groupshared float3 gs_cacheLighting[36];
groupshared float gs_cacheDepth[36];

float _RayMarchingLowResPercentage;

void FillUpsampleDataLDS(uint groupIndex, uint2 groupOrigin)
{
    // Define which value we will be acessing with this worker thread
    int acessCoordX = groupIndex % 6;
    int acessCoordY = groupIndex / 6;

    // Everything we are accessing is in intermediate res (half rez).
    uint2 traceGroupOrigin = groupOrigin / 2;

    // The initial position of the access
    int2 originXY = traceGroupOrigin - int2(1, 1);

    // Compute the sample position
    int2 sampleCoord = int2(clamp(originXY.x + acessCoordX, 0, _HalfScreenSize.x - 1), clamp(originXY.y + acessCoordY, 0, _HalfScreenSize.y - 1));

    // Sample and store into the LDS
    gs_cacheLighting[groupIndex] = LOAD_TEXTURE2D_X(_LowResolutionTexture, sampleCoord).xyz;
    // As an input we are not using the depth pyramid, but the full resolution depth (so we need to make sure to read from there for the upsample aswell).
    gs_cacheDepth[groupIndex] = LOAD_TEXTURE2D_X(_DepthTexture, sampleCoord * 2).x;
}

uint OffsetToLDSAdress(uint2 groupThreadId, int2 offset)
{
    // Compute the tap coordinate in the 6x6 grid
    uint2 tapAddress = (uint2)((int2)(groupThreadId / 2 + 1) + offset);
    return clamp((uint)(tapAddress.x) % 6 + tapAddress.y * 6, 0, 35);
}

// Function that fills the struct as we cannot use arrays
void FillUpsampleNeighborhoodData_2x2(int2 groupThreadId, int subRegionIdx, out NeighborhoodUpsampleData2x2_RGB neighborhoodData)
{
    // Fill the sample data
    int tapIdx = OffsetToLDSAdress(groupThreadId, int2((int)_TapOffsets[2 * subRegionIdx].x, (int)_TapOffsets[2 * subRegionIdx].y));
    neighborhoodData.lowValue0 = max(0, (gs_cacheLighting[tapIdx]));
    neighborhoodData.lowDepth.x = gs_cacheDepth[tapIdx];
    neighborhoodData.lowWeight.x = _DistanceBasedWeights[subRegionIdx].x;

    tapIdx = OffsetToLDSAdress(groupThreadId, int2((int)_TapOffsets[2 * subRegionIdx].z, (int)_TapOffsets[2 * subRegionIdx].w));
    neighborhoodData.lowValue1 = max(0, (gs_cacheLighting[tapIdx]));
    neighborhoodData.lowDepth.y = gs_cacheDepth[tapIdx];
    neighborhoodData.lowWeight.y = _DistanceBasedWeights[subRegionIdx].y;

    tapIdx = OffsetToLDSAdress(groupThreadId, int2((int)_TapOffsets[2 * subRegionIdx + 1].x, (int)_TapOffsets[2 * subRegionIdx + 1].y));
    neighborhoodData.lowValue2 = max(0, (gs_cacheLighting[tapIdx]));
    neighborhoodData.lowDepth.z = gs_cacheDepth[tapIdx];
    neighborhoodData.lowWeight.z = _DistanceBasedWeights[subRegionIdx].z;

    tapIdx = OffsetToLDSAdress(groupThreadId, int2((int)_TapOffsets[2 * subRegionIdx + 1].z, (int)_TapOffsets[2 * subRegionIdx + 1].w));
    neighborhoodData.lowValue3 = max(0, (gs_cacheLighting[tapIdx]));
    neighborhoodData.lowDepth.w = gs_cacheDepth[tapIdx];
    neighborhoodData.lowWeight.w = _DistanceBasedWeights[subRegionIdx].w;
}

// The output of our upscaling pass
RW_TEXTURE2D_X(float3, _OutputUpscaledTexture);

[numthreads(BILATERAL_UPSAMPLE_TILE_SIZE, BILATERAL_UPSAMPLE_TILE_SIZE, 1)]
void BilateralUpSampleColorHalf(uint3 currentCoord : SV_DispatchThreadID,
                    int groupIndex : SV_GroupIndex,
                    uint2 groupThreadId : SV_GroupThreadID,
                    uint2 groupId : SV_GroupID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(currentCoord.z);

    // Only 36 workers of the 64 region do the pre-fetching
    if (groupIndex < 36)
    {
        // Load 1 value per thread
        FillUpsampleDataLDS(groupIndex, groupId * 8);
    }

    // Make sure all values are loaded in LDS by now.
    GroupMemoryBarrierWithGroupSync();

    // If out of bounds, discard
    if (any(currentCoord.xy > uint2(_ScreenSize.xy)))
        return;

    // Read the depth value as early as possible and use it as late as possible
    float hiResDepth = LOAD_TEXTURE2D_X(_DepthTexture, currentCoord.xy).x;

    // Tap the neighborhood data from
    NeighborhoodUpsampleData2x2_RGB upsampleData;
    int localIndex = (currentCoord.x & 1) + (currentCoord.y & 1) * 2;
    FillUpsampleNeighborhoodData_2x2(groupThreadId, localIndex, upsampleData);

    // Upscale and return the result
    _OutputUpscaledTexture[COORD_TEXTURE2D_X(currentCoord.xy)] = BilUpColor2x2_RGB(hiResDepth, upsampleData);
}

SAMPLER(sampler_LinearClamp);
SAMPLER(sampler_PointClamp);

[numthreads(BILATERAL_UPSAMPLE_TILE_SIZE, BILATERAL_UPSAMPLE_TILE_SIZE, 1)]
void BilateralUpSampleColor(uint3 currentCoord : SV_DispatchThreadID)
{
    UNITY_XR_ASSIGN_VIEW_INDEX(currentCoord.z);

    // Only 36 workers of the 64 region do the pre-fetching
    // If out of bounds, discard
    if (any(currentCoord.xy >= uint2(_ScreenSize.xy)))
        return;

    // Read the depth value as early as possible and use it as late as possible
    // TODO: do a gather here.
    float4 depthNeighborhood = float4(
        LOAD_TEXTURE2D_X(_DepthTexture, currentCoord.xy + uint2(0,1)).x,
        LOAD_TEXTURE2D_X(_DepthTexture, currentCoord.xy + uint2(1,1)).x,
        LOAD_TEXTURE2D_X(_DepthTexture, currentCoord.xy + uint2(1,0)).x,
        LOAD_TEXTURE2D_X(_DepthTexture, currentCoord.xy + uint2(0,0)).x);

    float closestDepth = max(depthNeighborhood.x, max(depthNeighborhood.y, max(depthNeighborhood.z, depthNeighborhood.w)));

    float2 uvAtLowRes = min((currentCoord.xy) * _RayMarchingLowResPercentage + 0.5, _HalfScreenSize.xy - 1) * _ScreenSize.zw;
    float2 sampleUV = ClampAndScaleUVForBilinear(uvAtLowRes); 

    float2 samplePixel = sampleUV * _ScreenSize.xy;
    float2 bottomRight = frac(samplePixel + 0.5);
    float2 topLeft = 1.0 - bottomRight;
    float4 linearWeights = float4(
                    topLeft.x * bottomRight.y,
                    bottomRight.x * bottomRight.y,
                    bottomRight.x * topLeft.y,
                    topLeft.x * topLeft.y);

    float4 reds = GATHER_RED_TEXTURE2D_X(_LowResolutionTexture, sampler_LinearClamp, sampleUV);
    float4 greens = GATHER_GREEN_TEXTURE2D_X(_LowResolutionTexture, sampler_LinearClamp, sampleUV);
    float4 blues = GATHER_BLUE_TEXTURE2D_X(_LowResolutionTexture, sampler_LinearClamp, sampleUV);

    float3 trueCol = SAMPLE_TEXTURE2D_X_LOD(_LowResolutionTexture, sampler_LinearClamp, sampleUV, 0.0).rgb;
    float3 s0  = float3(reds.x, greens.x, blues.x);
    float3 s1  = float3(reds.y, greens.y, blues.y);
    float3 s2  = float3(reds.z, greens.z, blues.z);
    float3 s3 = float3(reds.w, greens.w, blues.w);

    float3 bilCol = BilUpColor3WithWeight(closestDepth, depthNeighborhood, s0, s1, s2, s3, linearWeights);

    // Upscale and return the result
    _OutputUpscaledTexture[COORD_TEXTURE2D_X(currentCoord.xy)] = bilCol;
}