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

#pragma kernel ResolveFlatNormals

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

Texture2D<float> _ScreenDepths;
RWTexture2D<float3> _ScreenFlatNormals;

float4x4 _ClipToWorldTransform;

struct PixelInfo
{
    float2 uvPos;
    float depthNDC;
};

PixelInfo GetPixelInfo(uint2 pixelPos, float2 rcpScreenSize)
{
    PixelInfo texCoordAndDepth;
    texCoordAndDepth.depthNDC = LoadNdcDepth(_ScreenDepths, pixelPos);
    texCoordAndDepth.uvPos = PixelPosToUvPos(pixelPos, rcpScreenSize);
    return texCoordAndDepth;
}

struct EdgeInfo
{
    PixelInfo nearPixelInfo;
    float extrapolatedDepthNDC;
};

EdgeInfo ComputeEdgeInfo(uint2 centerPixelPos, int2 deltaPixelPos, float2 reciprocalScreenSize)
{
    EdgeInfo edgeInfo;
    edgeInfo.nearPixelInfo = GetPixelInfo(centerPixelPos + deltaPixelPos, reciprocalScreenSize);
    const float farPixelDepthNDC = LoadNdcDepth(_ScreenDepths, centerPixelPos + 2*deltaPixelPos);
    // This works for both perspective and orthographic because in perspective projection
    // to take into account perspective interpolation, we should interpolate 1/LinearZ and take the reciprocal of that.
    // But NDCDepth is already proportional to 1/LinearZ.
    //
    // To extrapolate the segment [far, near] beyond near with the same spacing as between near and far, we want to
    // return near + (near - far) = 2 * near - far.
    edgeInfo.extrapolatedDepthNDC = 2 * edgeInfo.nearPixelInfo.depthNDC - farPixelDepthNDC;
    return edgeInfo;
}

float3 GetNeighborWorldPos(uint2 centerPixelPos, float centerDepthNDC, int2 deltaPixelPos, float2 reciprocalScreenSize, inout bool handednessFlip)
{
    // Compute the edge passing through the two neighboring pixels on the positive side
    const EdgeInfo edgeInfoPositiveSide = ComputeEdgeInfo(centerPixelPos, deltaPixelPos, reciprocalScreenSize);

    // Compute the edge passing through the two neighboring pixels on the negative side
    const EdgeInfo edgeInfoNegativeSide = ComputeEdgeInfo(centerPixelPos, -deltaPixelPos, reciprocalScreenSize);

    PixelInfo neighborPixelInfo;
    if (abs(centerDepthNDC - edgeInfoNegativeSide.extrapolatedDepthNDC) < abs(centerDepthNDC - edgeInfoPositiveSide.extrapolatedDepthNDC))
    {
        neighborPixelInfo = edgeInfoNegativeSide.nearPixelInfo;
        handednessFlip = !handednessFlip;
    }
    else
    {
        neighborPixelInfo = edgeInfoPositiveSide.nearPixelInfo;
    }

    return ComputeWorldSpacePosition(neighborPixelInfo.uvPos, neighborPixelInfo.depthNDC, _ClipToWorldTransform);
}

[numthreads(8, 8, 1)]
void ResolveFlatNormals(uint2 centerPixelPos : SV_DispatchThreadID)
{
    // This is inspired by https://atyuwen.github.io/posts/normal-reconstruction/
    // In addition to the center sample, we take 8 depth samples in a cross around the center.
    //        2
    //        3
    //  0  1  x  5  4
    //        7
    //        6
    // Of these 8, we pick the best in the horizontal direction and the best in the vertical direction
    // and use these to approximate the triangle surface orientation. By the best, this means picking
    // the point for which the extrapolation of the segment they lie on towards the central pixel
    // minimizes the difference in depth with that central pixel's depth.
    // Note: this technique could be improved for some corner (pun intended as the algorithm sometimes
    // doesn't return the correct true triangle normal in some tight corners) cases with some extra
    // diagonal extrapolation.

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

    if (any(screenSize <= centerPixelPos))
        return;

    const float centerNdcDepth = LoadNdcDepth(_ScreenDepths, centerPixelPos);
    if (centerNdcDepth == invalidNdcDepth)
        return;

    const float2 reciprocalScreenSize = rcp(float2(screenSize));
    const float2 centerUVPos = PixelPosToUvPos(centerPixelPos, reciprocalScreenSize);
    const float3 centerWorldPos = ComputeWorldSpacePosition(centerUVPos, centerNdcDepth, _ClipToWorldTransform);

    bool handednessFlip = false; // Account for handedness and winding order.

    const float3 horizontalWorldPos = GetNeighborWorldPos(centerPixelPos, centerNdcDepth, int2(1, 0), reciprocalScreenSize, handednessFlip);
    const float3 ddxPos = horizontalWorldPos - centerWorldPos;

    const float3 verticalWorldPos = GetNeighborWorldPos(centerPixelPos, centerNdcDepth, -int2(0, 1), reciprocalScreenSize, handednessFlip);
    const float3 ddyPos = verticalWorldPos - centerWorldPos;

    float3 flatNormal = normalize(cross(ddxPos, ddyPos));
    if (handednessFlip)
        flatNormal = -flatNormal;

    _ScreenFlatNormals[centerPixelPos] = flatNormal;
}