// Crest Water System
// Copyright © 2024 Wave Harmonic. All rights reserved.

#pragma kernel CrestQueryDisplacement d_CrestDisplacement
#pragma kernel CrestQueryFlow d_CrestFlow
#pragma kernel CrestQueryDepth d_CrestDepth

// Must match value in script.
#define GROUP_SIZE 64

StructuredBuffer<float3> _Crest_QueryPositions_MinimumGridSizes;
RWStructuredBuffer<float3> _Crest_Target;

#include "HLSLSupport.cginc"

#include "Packages/com.waveharmonic.crest/Runtime/Shaders/Library/Macros.hlsl"
#include "Packages/com.waveharmonic.crest/Runtime/Shaders/Library/Globals.hlsl"
#include "Packages/com.waveharmonic.crest/Runtime/Shaders/Library/InputsDriven.hlsl"
#include "Packages/com.waveharmonic.crest/Runtime/Shaders/Library/Helpers.hlsl"
#include "Packages/com.waveharmonic.crest/Runtime/Shaders/Library/Cascade.hlsl"

#if d_CrestDisplacement
#define d_CrestSample SampleDisplacement
#define d_CrestMake MakeAnimatedWaves
#define d_CrestComponents float3
#define d_CrestReturn result
#elif d_CrestDepth
#define d_CrestSample SampleSignedDepthFromSeaLevelAndDistance
#define d_CrestMake MakeDepth
#define d_CrestComponents float2
#define d_CrestReturn float3(result.x, result.y, 0.0)
#else
#define d_CrestSample SampleFlow
#define d_CrestMake MakeFlow
#define d_CrestComponents float2
// Unfortunately we don't support float2 vs float3s yet, on the C# side
#define d_CrestReturn float3(result.x, 0.0, result.y)
#endif

m_CrestNameSpace


float3 Compute(const float2 i_Position, const float i_MinimumSlice, const float i_BaseScale)
{
    // Do not use last slice - this is a 'transition' slice used to cross-fade waves
    // between LOD resolutions to avoid pops. That being said, this will have clamped
    // samples leading to objects floating on waves that do not exist.
    uint slice0, slice1; float alpha;
    PosToSliceIndices(i_Position, i_MinimumSlice, g_Crest_LodCount - 2.0, i_BaseScale, slice0, slice1, alpha);

    const Cascade cascade0 = Cascade::d_CrestMake(slice0);
    const Cascade cascade1 = Cascade::d_CrestMake(slice1);

    const float weight0 = (1.0 - alpha) * Cascade::Make(slice0)._Weight;
    const float weight1 = (1.0 - weight0) * Cascade::Make(slice1)._Weight;

    d_CrestComponents result =
        weight0 * cascade0.d_CrestSample(i_Position) +
        weight1 * cascade1.d_CrestSample(i_Position);

    return d_CrestReturn;
}

void Query(const uint3 id)
{
    const float3 data = _Crest_QueryPositions_MinimumGridSizes[id.x];
    const float minimumGridSize = data.z;
    float2 position = data.xy;

    const float gridSizeSlice0 = Cascade::d_CrestMake(0)._Texel;
    // Displacements should not utilize the last slice which is used for transitioning
    // waves between sampling resolutions. While it might be ok to use the last slice
    // for other targets, we avoid using it to be consistent with displacements.
    const float minimumSlice = clamp(floor(log2(max(minimumGridSize / gridSizeSlice0, 1.0))), 0.0, g_Crest_LodCount - 2.0);

#if d_CrestDisplacement
    // Perform iteration to invert the displacement vector field - find position that displaces to query position,
    // and return displacement at that point.
    float2 undisplaced = position;
    for (int i = 0; i < 4; i++)
    {
        const float3 displacement = Compute(undisplaced, minimumSlice, g_Crest_WaterScale);
        const float2 error = (undisplaced + displacement.xz) - position;
        undisplaced -= error;
    }

    position = undisplaced;
#endif

    _Crest_Target[id.x] = Compute(position, minimumSlice, g_Crest_WaterScale);
}

m_CrestNameSpaceEnd

[numthreads(GROUP_SIZE, 1, 1)]
m_CrestKernelVariant(Query, Displacement)

[numthreads(GROUP_SIZE, 1, 1)]
m_CrestKernelVariant(Query, Flow)

[numthreads(GROUP_SIZE, 1, 1)]
m_CrestKernelVariant(Query, Depth)