#pragma kernel ReprojectFoam
#pragma kernel AttenuateFoam

#pragma only_renderers d3d11 playstation xboxone xboxseries vulkan metal switch

// #pragma enable_d3d11_debug_symbols

// Required to be defined for some includes
#define WATER_SIMULATION

// SRP generic includes
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/NormalSurfaceGradient.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/ShaderLibrary/ShaderVariables.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Water/WaterSystemDef.cs.hlsl"
#include "Packages/com.unity.render-pipelines.high-definition/Runtime/Water/ShaderVariablesWater.cs.hlsl"

// EvaluateNormals UAVS
Texture2D<float2> _WaterFoamBuffer;
RWTexture2D<float2> _WaterFoamBufferRW;
float4 _PreviousFoamRegionScaleOffset;

bool PreviousLocationInsideRegion(float2 foamPrevUV)
{
    return all(foamPrevUV >= 0.0) && all(foamPrevUV <= 1.0);
}

float2 FoamUVToWorld(float2 uv)
{
    uv = uv / _FoamRegionScale;

    float2 axis1 = _WaterForwardXZ;
    float2 axis2 = float2(-axis1.y, axis1.x);
    uv = uv.x * axis1 + uv.y * axis2;

    return uv + _FoamRegionOffset;
}

float2 GetPreviousFoamUV(float2 transformedPositionAWS)
{
    float2 deformationPosition = transformedPositionAWS - _PreviousFoamRegionScaleOffset.zw;

    float2 axis1 = _WaterForwardXZ;
    float2 axis2 = float2(-axis1.y, axis1.x);
    deformationPosition = float2(dot(deformationPosition, axis1), dot(deformationPosition, axis2));

    return deformationPosition * _PreviousFoamRegionScaleOffset.xy;
}

[numthreads(8, 8, 1)]
void ReprojectFoam(uint3 currentThread : SV_DispatchThreadID,
                        int groupIndex : SV_GroupIndex,
                        uint2 groupId : SV_GroupID,
                        uint2 groupThreadId : SV_GroupThreadID)
{
    // Extract the information about the pixel to process
    uint2 coord = currentThread.xy;

    // Evaluate the UV coordinates of this pixel
    float2 foamRegionCoord = (coord + 0.5f) / _WaterFoamRegionResolution - 0.5f;

    // Evaluate the world space position of this pixel
    float2 foamPosWS = FoamUVToWorld(foamRegionCoord);

    // Evaluate the previous foam region
    float2 foamPrevUV = GetPreviousFoamUV(foamPosWS) + 0.5f;

    // Output the normal and foam
    _WaterFoamBufferRW[coord] = PreviousLocationInsideRegion(foamPrevUV) ? SAMPLE_TEXTURE2D_LOD(_WaterFoamBuffer, s_linear_clamp_sampler, foamPrevUV, 0).xy : 0.0;
}


[numthreads(8, 8, 1)]
void AttenuateFoam(uint3 currentThread : SV_DispatchThreadID,
                        int groupIndex : SV_GroupIndex,
                        uint2 groupId : SV_GroupID,
                        uint2 groupThreadId : SV_GroupThreadID)
{
    // Extract the information about the pixel to process
    uint2 coord = currentThread.xy;

    // Output the normal and foam
    _WaterFoamBufferRW[coord] = saturate(LOAD_TEXTURE2D_LOD(_WaterFoamBuffer, coord, 0).xy);
}