#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal glcore ps5 switch2

#pragma kernel Visualize

#include "Packages/com.unity.render-pipelines.universal/ShaderLibrary/Core.hlsl"
#include "Packages/com.unity.render-pipelines.universal/ShaderLibrary/GBufferCommon.hlsl"
#include "Packages/com.unity.render-pipelines.core/Runtime/Sampling/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/Runtime/Sampling/Hashes.hlsl"
#include "Packages/com.unity.render-pipelines.core/Runtime/Lighting/SurfaceCache/PatchUtil.hlsl"
#include "Packages/com.unity.render-pipelines.core/Runtime/Lighting/SurfaceCache/RingBuffer.hlsl"

StructuredBuffer<uint> _RingConfigBuffer;
StructuredBuffer<SphericalHarmonics::RGBL1> _PatchIrradiances;
StructuredBuffer<uint> _PatchCellIndices;
StructuredBuffer<PatchUtil::PatchGeometry> _PatchGeometries;
RWStructuredBuffer<PatchUtil::PatchStatisticsSet> _PatchStatistics;
StructuredBuffer<uint> _CellPatchIndices;
StructuredBuffer<int3> _VolumeCascadeOffsets;

RWTexture2D<float4> _Result;
Texture2D<float> _ScreenDepths;
Texture2D<float3> _ScreenShadedNormals;
Texture2D<float3> _ScreenFlatNormals;

uint _RingConfigOffset;
uint _VolumeSpatialResolution;
uint _VolumeCascadeCount;
uint _ViewMode;
uint _ShowSamplePosition;
float4x4 _ClipToWorldTransform;
uint _FrameIdx;
float _VolumeVoxelMinSize;
float3 _VolumeTargetPos;

uint _DummyBlackbox; // By using a uniform we force the compiler to not make any assumptions.
float3 IsActuallyNan(float3 val) {
    return isnan(asfloat(_DummyBlackbox ^ asuint(val)));
}

float3 IsActuallyInf(float3 val) {
    return isinf(asfloat(_DummyBlackbox ^ asuint(val)));
}

[numthreads(8,8,1)]
void Visualize(uint2 threadIdx : SV_DispatchThreadID)
{
    const float3 green = float3(0, 1, 0);
    const float3 red = float3(1, 0, 0);
    const float3 pink = float3(1, 0, 1);
    const float3 white = float3(1, 1, 1);
    const float3 yellow = float3(1, 1, 0);

    uint2 screenSize;
    _ScreenDepths.GetDimensions(screenSize.x, screenSize.y);

    float3 output;

    const float pixelPosYPercentage = float(threadIdx.y) / screenSize.y;
    const float statusBarHeightPercentage = 0.05f;
    if (pixelPosYPercentage < statusBarHeightPercentage)
    {
        const RingBuffer::Config ringConfig = RingBuffer::LoadConfig(_RingConfigBuffer, _RingConfigOffset);
        const uint patchIdx = float(threadIdx.x) / screenSize.x * patchCapacity;
        if (RingBuffer::IsPositionInUse(ringConfig, patchIdx))
        {
            if (_PatchCellIndices[patchIdx] == PatchUtil::invalidCellIndex)
            {
                output = red;
            }
            else
            {
                output = PatchUtil::GetRank(_PatchStatistics[patchIdx].counters) == 0 ? green : yellow;
                if (pixelPosYPercentage < statusBarHeightPercentage * 0.5f)
                {
                    PatchUtil::PatchStatisticsSet stats = _PatchStatistics[patchIdx];
                    const uint frameSinceLastAccess = PatchUtil::GetFramesSinceLastAccess(_FrameIdx, PatchUtil::GetLastAccessFrame(stats.counters));
                    output *= 1.0f - float(frameSinceLastAccess) / float(PatchUtil::evictionThreshold);
                }
            }
        }
        else
        {
            output = white;
        }
    }
    else
    {
        const float ndcDepth = LoadNdcDepth(_ScreenDepths, threadIdx);
        if (ndcDepth == invalidNdcDepth)
            return;

        const float2 uv = (float2(threadIdx.xy) + 0.5f) / float2(screenSize);
        const float3 worldPos = ComputeWorldSpacePosition(uv, ndcDepth, _ClipToWorldTransform);
        const float3 worldShadedNormal = UnpackGBufferNormal(_ScreenShadedNormals[threadIdx.xy]);
        const float3 worldFlatNormal = _ScreenFlatNormals[threadIdx.xy];

        PatchUtil::VolumeParamSet volumeParams;
        volumeParams.targetPos = _VolumeTargetPos;
        volumeParams.spatialResolution = _VolumeSpatialResolution;
        volumeParams.voxelMinSize = _VolumeVoxelMinSize;
        volumeParams.cascadeOffsets = _VolumeCascadeOffsets;
        volumeParams.cascadeCount = _VolumeCascadeCount;

        PatchUtil::VolumePositionResolution posResolution = PatchUtil::ResolveVolumePosition(worldPos, volumeParams);
        output = float3(1.0f, 0.0f, 1.0f);
        if (posResolution.isValid())
        {
            const uint directionIdx = PatchUtil::GetDirectionIndex(worldFlatNormal, PatchUtil::volumeAngularResolution);
            const uint3 positionStorageSpace = PatchUtil::ConvertVolumeSpaceToStorageSpace(posResolution.positionVolumeSpace, _VolumeSpatialResolution, _VolumeCascadeOffsets[posResolution.cascadeIdx]);
            const uint cellIdx = PatchUtil::GetCellIndex(posResolution.cascadeIdx, positionStorageSpace, directionIdx, _VolumeSpatialResolution, PatchUtil::volumeAngularResolution);
            const uint patchIdx = _CellPatchIndices[cellIdx];
            const bool hasPatch = (patchIdx != PatchUtil::invalidPatchIndex);

            SphericalHarmonics::RGBL1 patchIrradiance = (SphericalHarmonics::RGBL1)0;
            float3 patchPos = 0; // initializing value only to silence shader warning
            if (hasPatch)
            {
                patchIrradiance = _PatchIrradiances[patchIdx];
                patchPos = _PatchGeometries[patchIdx].position;
                PatchUtil::WriteLastFrameAccess(_PatchStatistics, patchIdx, _FrameIdx);
            }

            if (_ShowSamplePosition == 1 && hasPatch && length(worldPos - patchPos) < 0.1f * PatchUtil::GetVoxelSize(_VolumeVoxelMinSize, posResolution.cascadeIdx))
            {
                output = 30.0f;
            }
            else if (_ViewMode == 0)
            {
                output = float3(
                    UintToFloat01(LowBiasHash32(cellIdx, 0)),
                    UintToFloat01(LowBiasHash32(cellIdx, 1)),
                    UintToFloat01(LowBiasHash32(cellIdx, 2)));
            }
            else if (_ViewMode == 1)
            {
                if (hasPatch)
                {
                    const float3 shEval = SphericalHarmonics::Eval(patchIrradiance, worldShadedNormal);
                    if (any(IsActuallyNan(shEval)))
                        output = red;
                    else if (any(IsActuallyInf(shEval)))
                        output = green;
                    else
                        output = max(0, shEval);
                }
            }
            else if (_ViewMode == 2) // fast irradiance
            {
                if (hasPatch)
                {
                    output = _PatchStatistics[patchIdx].mean * SphericalHarmonics::y0;
                }
            }
            else if (_ViewMode == 3) // coefficient of variation
            {
                if (hasPatch)
                {
                    const PatchUtil::PatchStatisticsSet stats = _PatchStatistics[patchIdx];
                    output = sqrt(stats.variance) / max(0.01f, stats.mean) / 20.0f;
                }
            }
            else if (_ViewMode == 4)
            {
                if (hasPatch)
                {
                    const PatchUtil::PatchStatisticsSet stats = _PatchStatistics[patchIdx];
                    float3 longIrradiance = _PatchIrradiances[patchIdx].l0;
                    float3 shortStdDev = sqrt(stats.variance);
                    float3 drift = abs(longIrradiance - stats.mean) / max(0.001f, shortStdDev);
                    output = drift * 0.25f;
                }
            }
            else if (_ViewMode == 5)
            {
                if (hasPatch)
                {
                    float3 variance = _PatchStatistics[patchIdx].variance;
                    if (any(isinf(variance)) || any(isnan(variance)))
                        output = pink;
                    else if (any(variance < 0))
                        output = red; // Since this is an estimate, negative values may happen.
                    else
                        output = sqrt(variance) * 0.1f;
                }
            }
            else if (_ViewMode == 6)
            {
                if (hasPatch)
                {
                    const PatchUtil::PatchCounterSet counterSet = _PatchStatistics[patchIdx].counters;
                    const uint updateCount = PatchUtil::GetUpdateCount(counterSet);
                    const float s = float(updateCount) / PatchUtil::updateMax;
                    const float3 redGreenMix = lerp(red, green, s);
                    const float3 checker = UintToFloat01(LowBiasHash32(cellIdx, 0)); // To make cells visible.
                    output = lerp(redGreenMix, checker, 0.25f);
                }
            }
            else if (_ViewMode == 7)
            {
                output = worldFlatNormal * 0.5f + 0.5f;
            }
        }
    }

    _Result[threadIdx.xy] = float4(output, 1.0f);
}