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

#pragma kernel CrestInitialize
#pragma kernel CrestInitializeGroundHeight
#pragma kernel CrestAdvect
#pragma kernel CrestUpdateHeight
#pragma kernel CrestHeightOvershootReduction
#pragma kernel CrestUpdateVelocity
#pragma kernel CrestBlurHeight
#pragma kernel CrestBlur
#pragma kernel CrestMaskEdge
#pragma kernel CrestExpandEdge
#pragma kernel CrestFinalizeHeight
#pragma kernel CrestInjectShape
#pragma kernel CrestInjectLevel
#pragma kernel CrestInjectFoam
#pragma kernel CrestInjectFlow
#pragma kernel CrestBakeLevel
#pragma kernel CrestBakeFoam
#pragma kernel CrestBakeFlow

#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"

#define FLT_MIN 1.175494351e-38

SamplerState linear_clamp_sampler;

Texture2D<float>   _Crest_LevelSource;
Texture2D<float>   _Crest_HeightSource;
RWTexture2D<float> _Crest_HeightTarget;
Texture2D<float>   _Crest_VelocityXSource;
RWTexture2D<float> _Crest_VelocityXTarget;
Texture2D<float>   _Crest_VelocityYSource;
RWTexture2D<float> _Crest_VelocityYTarget;
Texture2D<float>   _Crest_GroundHeightSource;
RWTexture2D<float> _Crest_GroundHeightTarget;
Texture2D<float>   _Crest_MaskSource;
RWTexture2D<float> _Crest_MaskTarget;

// Injector Targets.
RWTexture2DArray<float>  _Crest_FoamTarget;
RWTexture2DArray<float2> _Crest_FlowTarget;
RWTexture2DArray<float4> _Crest_ShapeTarget;
RWTexture2DArray<float> _Crest_LevelTarget;

// Bake Targets.
RWTexture2D<float>  _Crest_FoamBakeTarget;
RWTexture2D<float2> _Crest_FlowBakeTarget;
RWTexture2D<float>  _Crest_LevelBakeTarget;

// Depth Probe.
Texture2D<float2> _Crest_DepthProbe;

CBUFFER_START(CrestBuffer)
float _Crest_SimDeltaTime;

float _Crest_SimulationDeltaTime;
float _Crest_DomainWidth;
float3 _Crest_DomainOrigin;
uint _Crest_Resolution;
float _Crest_DrainWaterAtBoundaries;
float _Crest_Evaporation;
float _Crest_Friction;
float _Crest_MaximumVelocity;
float _Crest_AddAdditionalWater;
float _Crest_TexelSize;

bool _Crest_MacCormackAdvection;
bool _Crest_MacCormackAdvectionForHeight;
bool _Crest_UpwindHeight;
bool _Crest_DepthLimiter;
bool _Crest_InjectLevel;

float _Crest_OvershootReductionStrength;

float _Crest_ShallowMinimumDepth;
float _Crest_ShallowMaximumDepth;
float _Crest_BlendPushUpStrength;

float _Crest_InjectionStrength;
float _Crest_InjectionMaximum;
float _Crest_Weight;

bool _Crest_SampleFromDepthProbe;
bool _Crest_SampleFromLevelInputs;
float _Crest_DepthProbeHeightOffset;
float _Crest_DepthProbeResolution;

float _Crest_MarginWidth;
CBUFFER_END

m_CrestNameSpace

float2 IdToWorldXZ_H(uint2 id)
{
    // H grid is centered
    return ((id + 0.5) / _Crest_Resolution - 0.5) * _Crest_DomainWidth + _Crest_DomainOrigin.xz;
}
float2 IdToWorldXZ_GroundHeight(uint2 id)
{
    // H and GroundHeight are aligned
    return IdToWorldXZ_H(id);
}
float2 IdToWorldXZ_Vx(uint2 id)
{
    // Vx grid is offset right half a texel
    return ((id + float2(1.0, 0.5)) / _Crest_Resolution - 0.5) * _Crest_DomainWidth + _Crest_DomainOrigin.xz;
}
float2 IdToWorldXZ_Vy(uint2 id)
{
    // Vy grid is offset forward half a texel
    return ((id + float2(0.5, 1.0)) / _Crest_Resolution - 0.5) * _Crest_DomainWidth + _Crest_DomainOrigin.xz;
}

float2 WorldXZToUv_H(float2 worldXZ)
{
    // H grid is centered
    return (worldXZ - _Crest_DomainOrigin.xz) / _Crest_DomainWidth + 0.5;
}
float2 WorldXZToUv_GroundHeight(float2 worldXZ)
{
    // Aligned to H
    return WorldXZToUv_H(worldXZ);
}
float2 WorldXZToUv_Vx(float2 worldXZ)
{
    // Vx grid is offset right half a texel
    return (worldXZ - _Crest_DomainOrigin.xz) / _Crest_DomainWidth + 0.5 - float2(0.5 / _Crest_Resolution, 0.0);
}
float2 WorldXZToUv_Vy(float2 worldXZ)
{
    // Vy grid is offset forward half a texel
    return (worldXZ - _Crest_DomainOrigin.xz) / _Crest_DomainWidth + 0.5 - float2(0.0, 0.5 / _Crest_Resolution);
}

float SampleWaterLevel(float2 worldXZ)
{
    return _Crest_LevelSource.SampleLevel(linear_clamp_sampler, WorldXZToUv_H(worldXZ), 0.0).x;
}
float SampleH(float2 worldXZ)
{
    return _Crest_HeightSource.SampleLevel(linear_clamp_sampler, WorldXZToUv_H(worldXZ), 0.0).x;
}
float SampleGroundHeight(float2 worldXZ)
{
    return _Crest_GroundHeightSource.SampleLevel(linear_clamp_sampler, WorldXZToUv_GroundHeight(worldXZ), 0.0).x;
}
float SampleVx0(float2 worldXZ)
{
    return _Crest_VelocityXSource.SampleLevel(linear_clamp_sampler, WorldXZToUv_Vx(worldXZ), 0.0).x;
}
float SampleVy0(float2 worldXZ)
{
    return _Crest_VelocityYSource.SampleLevel(linear_clamp_sampler, WorldXZToUv_Vy(worldXZ), 0.0).x;
}


//
// Initialize
//

void Initialize(const uint3 i_ID)
{
    const uint2 id = i_ID.xy;

    const float2 worldXZ = IdToWorldXZ_H(id);

    uint slice0, slice1; float alpha;
    PosToSliceIndices(worldXZ, 0.0, g_Crest_LodCount - 1.0, g_Crest_WaterScale, slice0, slice1, alpha);

    float h = g_Crest_WaterCenter.y;
    const float weight0 = (1.0 - alpha) * g_Crest_CascadeData[slice0].y;
    const float weight1 = (1.0 - weight0) * g_Crest_CascadeData[slice1].y;

    if (_Crest_InjectLevel)
    {
        if (_Crest_SampleFromLevelInputs)
        {
            h += SampleWaterLevel(worldXZ);
        }
        else
        {
            h +=
                weight0 * Cascade::MakeLevel(slice0).SampleLevel(worldXZ) +
                weight1 * Cascade::MakeLevel(slice1).SampleLevel(worldXZ);
        }
    }

    h -= _Crest_DomainOrigin.y;

    h = max(0.0, h - SampleGroundHeight(worldXZ));

    h += _Crest_AddAdditionalWater;

    _Crest_HeightTarget[id] = h;
    _Crest_VelocityXTarget[id] = 0.0;
    _Crest_VelocityYTarget[id] = 0.0;
}

// Called first.
void InitializeGroundHeight(const uint3 i_ID)
{
    const uint2 id = i_ID.xy;

    const float2 worldXZ = IdToWorldXZ_GroundHeight(id);

    uint slice0, slice1; float alpha;
    PosToSliceIndices(worldXZ, 0.0, g_Crest_LodCount - 1.0, g_Crest_WaterScale, slice0, slice1, alpha);

    const float weight0 = (1.0 - alpha) * g_Crest_CascadeData[slice0].y;
    const float weight1 = (1.0 - weight0) * g_Crest_CascadeData[slice1].y;

    float g;
    if (_Crest_SampleFromDepthProbe)
    {
        // Sample from the Depth Probe directly otherwise we bake in LOD losses.
        g = _Crest_DepthProbe.SampleLevel(linear_clamp_sampler, id.xy / (float)_Crest_Resolution, 0.0).x;
        g += _Crest_DepthProbeHeightOffset;
    }
    else
    {
        g = weight0 * Cascade::MakeDepth(slice0).SampleSceneHeight(worldXZ) +
            weight1 * Cascade::MakeDepth(slice1).SampleSceneHeight(worldXZ);
    }

    // LOD losses for height should be acceptable providing height is flat (which it should be).
    // Bake water level into ground by offsetting ground with water level.
    if (!_Crest_InjectLevel)
    {
        g -=
            weight0 * Cascade::MakeLevel(slice0).SampleLevel(worldXZ) +
            weight1 * Cascade::MakeLevel(slice1).SampleLevel(worldXZ);
    }

    // Mask is directly based off depth at the moment.
    float depth = g_Crest_WaterCenter.y - g;
    _Crest_MaskTarget[id] = 1.0 - saturate((depth - _Crest_ShallowMinimumDepth) / _Crest_ShallowMaximumDepth);

    g = max(0.0, g - _Crest_DomainOrigin.y);

    _Crest_GroundHeightTarget[id] = g;
}


//
// Advect
//

void Advect(const uint3 id)
{
    // Vx
    if (id.z == 0)
    {
        const float2 worldXZ = IdToWorldXZ_Vx(id.xy);

        // Advect forwards to predict advected position
        const float2 worldXZAdvected_prediction = worldXZ - _Crest_SimulationDeltaTime * float2(SampleVx0(worldXZ), SampleVy0(worldXZ));

        float2 worldXZAdvected = worldXZAdvected_prediction;
        const float vxSemiLagrangian = SampleVx0(worldXZAdvected);
        float vxAdvected = vxSemiLagrangian;
        if (_Crest_MacCormackAdvection)
        {
            // Update prediction by advecting backwards to obtain a starting position and using the difference as an error term

            // Advect forwards to predict original position
            const float2 worldXZ_prediction = worldXZAdvected_prediction + _Crest_SimulationDeltaTime * float2(vxSemiLagrangian, SampleVy0(worldXZAdvected_prediction));

            // Remove half of prediction error to yield average of both
            worldXZAdvected -= (worldXZ_prediction - worldXZ) / 2.0;

            float vxAdvectedNew = SampleVx0(worldXZAdvected);

            const float4 vxsUsedForFirstSLStep = _Crest_VelocityXSource.Gather(linear_clamp_sampler, WorldXZToUv_Vx(worldXZ));
            const float vxMin = min(min(vxsUsedForFirstSLStep[0], vxsUsedForFirstSLStep[1]), min(vxsUsedForFirstSLStep[2], vxsUsedForFirstSLStep[3]));
            const float vxMax = max(max(vxsUsedForFirstSLStep[0], vxsUsedForFirstSLStep[1]), max(vxsUsedForFirstSLStep[2], vxsUsedForFirstSLStep[3]));

            if (vxAdvectedNew >= vxMin && vxAdvectedNew <= vxMax)
            {
                vxAdvected = vxAdvectedNew;
            }
        }

        // Abs to silence 3571 which for some reason cannot disable with pragma even
        // HLSLSupport can do exactly that.
        const float h_vx = SampleH(worldXZAdvected);
        const float friction = _Crest_Friction * _Crest_SimulationDeltaTime / max(0.001f, pow(abs(h_vx), 1.333333f));
        vxAdvected -= abs(vxAdvected) * vxAdvected * friction;

        vxAdvected = clamp(vxAdvected, -_Crest_MaximumVelocity, _Crest_MaximumVelocity);

        _Crest_VelocityXTarget[id.xy] = vxAdvected;
    }

    // Vy
    if (id.z == 1)
    {
        const float2 worldXZ = IdToWorldXZ_Vy(id.xy);

        // Advect forwards to predict advected position
        const float2 worldXZAdvected_prediction = worldXZ - _Crest_SimulationDeltaTime * float2(SampleVx0(worldXZ), SampleVy0(worldXZ));

        float2 worldXZAdvected = worldXZAdvected_prediction;
        const float vySemiLagrangian = SampleVy0(worldXZAdvected);
        float vyAdvected = vySemiLagrangian;
        if (_Crest_MacCormackAdvection)
        {
            // Update prediction by advecting backwards to obtain a starting position and using the difference as an error term

            // Advect forwards to predict original position
            const float2 worldXZ_prediction = worldXZAdvected_prediction + _Crest_SimulationDeltaTime * float2(SampleVx0(worldXZAdvected_prediction), vySemiLagrangian);

            // Remove half of prediction error to yield average of both
            worldXZAdvected -= (worldXZ_prediction - worldXZ) / 2.0;

            float vyAdvectedNew = SampleVy0(worldXZAdvected);

            const float4 vysUsedForFirstSLStep = _Crest_VelocityYSource.Gather(linear_clamp_sampler, WorldXZToUv_Vy(worldXZ));
            const float vyMin = min(min(vysUsedForFirstSLStep[0], vysUsedForFirstSLStep[1]), min(vysUsedForFirstSLStep[2], vysUsedForFirstSLStep[3]));
            const float vyMax = max(max(vysUsedForFirstSLStep[0], vysUsedForFirstSLStep[1]), max(vysUsedForFirstSLStep[2], vysUsedForFirstSLStep[3]));

            if (vyAdvectedNew >= vyMin && vyAdvectedNew <= vyMax)
            {
                vyAdvected = vyAdvectedNew;
            }
        }

        // Abs to silence 3571 which for some reason cannot disable with pragma even
        // HLSLSupport can do exactly that.
        const float h_vy = SampleH(worldXZAdvected);
        const float friction = _Crest_Friction * _Crest_SimulationDeltaTime / max(0.001f, pow(abs(h_vy), 1.333333f));
        vyAdvected -= abs(vyAdvected) * vyAdvected * friction;

        vyAdvected = clamp(vyAdvected, -_Crest_MaximumVelocity, _Crest_MaximumVelocity);

        _Crest_VelocityYTarget[id.xy] = vyAdvected;
    }

    // H
    if (id.z == 2)
    {
        const float2 worldXZ = IdToWorldXZ_H(id.xy);

        // Advect forwards to predict advected position
        const float2 worldXZAdvected_prediction = worldXZ - _Crest_SimulationDeltaTime * float2(SampleVx0(worldXZ), SampleVy0(worldXZ));

        float2 worldXZAdvected = worldXZAdvected_prediction;
        if (_Crest_MacCormackAdvectionForHeight)
        {
            // Update prediction by advecting backwards to obtain a starting position and using the difference as an error term

            // Advect forwards to predict original position
            const float2 worldXZ_prediction = worldXZAdvected_prediction + _Crest_SimulationDeltaTime * float2(SampleVx0(worldXZAdvected_prediction), SampleVy0(worldXZAdvected_prediction));

            // Remove half of prediction error to yield average of both
            worldXZAdvected -= (worldXZ_prediction - worldXZ) / 2.0;
        }

        _Crest_HeightTarget[id.xy] = SampleH(worldXZAdvected);
    }
}


//
// Update Height
//

float HeightAdjustment(uint2 id)
{
    const float gravity = 9.81f;
    float hn = _Crest_HeightSource[id + uint2(0, 1)];
    float hs = _Crest_HeightSource[id - uint2(0, 1)];
    float he = _Crest_HeightSource[id + uint2(1, 0)];
    float hw = _Crest_HeightSource[id - uint2(1, 0)];
    float h_avg = 0.25 * (hn + hs + he + hw);
    float beta = 2.0;
    float minThickness = beta * _Crest_TexelSize / (gravity * _Crest_SimulationDeltaTime);
    return max(0.0, h_avg - minThickness);
}

void UpdateHeight(const uint3 i_ID)
{
    const int2 id = i_ID.xy;

    // Fix height at border to prevent instability.
    if (i_ID.x == 0 || i_ID.y == 0 || i_ID.x == _Crest_Resolution - 1 || i_ID.y == _Crest_Resolution - 1)
    {
        _Crest_HeightTarget[id] = _Crest_HeightSource[id];
        return;
    }

    float h = _Crest_HeightSource[id];

    const uint x = id.x;
    const uint y = id.y;

    const float2 worldXZ = IdToWorldXZ_H(id);

    float pump = 0.0;

    if (h > 0.01)
    {
        // Drain water out at boundary of domain.
        {
            const float2 offset = abs(id / float(_Crest_Resolution) - 0.5);
            pump += _Crest_DrainWaterAtBoundaries * smoothstep(0.4, 0.5, max(offset.x, offset.y));
        }

        // Blend waves into simulation. Pumps water in.
        {
            uint slice0, slice1; float alpha;
            PosToSliceIndices(worldXZ, 0.0, g_Crest_LodCount - 1.0, g_Crest_WaterScale, slice0, slice1, alpha);

            const float weight0 = (1.0 - alpha) * g_Crest_CascadeData[slice0].y;
            const float weight1 = (1.0 - weight0) * g_Crest_CascadeData[slice1].y;

            float waterY = g_Crest_WaterCenter.y;

            // If injecting to height then ignore waves.
            if (_Crest_InjectLevel)
            {
                if (_Crest_SampleFromLevelInputs)
                {
                    waterY += SampleWaterLevel(worldXZ);
                }
                else
                {
                    waterY +=
                        weight0 * Cascade::MakeLevel(slice0).SampleLevel(worldXZ) +
                        weight1 * Cascade::MakeLevel(slice1).SampleLevel(worldXZ);
                }
            }
            else
            {
                waterY +=
                    weight0 * Cascade::MakeAnimatedWaves(slice0).SampleAnimatedWaves(worldXZ).y +
                    weight1 * Cascade::MakeAnimatedWaves(slice1).SampleAnimatedWaves(worldXZ).y;
            }

            float simulationY = h + _Crest_DomainOrigin.y + _Crest_GroundHeightSource[id];
            if (waterY > simulationY)
            {
                // This is tricksy. Don't apply forces from waves in deep, but also don't apply in very shallow.
                // The latter will mean that the SWS is held at sea level instead of being allowed to drop below.
                // I tried asserting disp.y > somevalue but that led to a whole different set of issues. Therefore
                // settle for a y=4x(1-x) weight that chooses the mid point of the mask blend. Seems to work.
                float maskWeight = 1.0;
                if (!_Crest_InjectLevel)
                {
                    maskWeight = _Crest_MaskSource[id];
                    maskWeight *= 4.0 * (1.0 - maskWeight);
                }

                pump += _Crest_BlendPushUpStrength * maskWeight * (waterY - simulationY);
            }
        }
    }

    // Update height based on total flow in minus total flow out.
    // Prevents wave reflections at domain boundary.
    float drain = 30.0;
    float vxp = (x == uint(_Crest_Resolution - 1)) ? drain : _Crest_VelocityXSource[id];
    float vyp = (y == uint(_Crest_Resolution - 1)) ? drain : _Crest_VelocityYSource[id];
    float vxm = (x == 0) ? -drain : _Crest_VelocityXSource[uint2(x - 1, y)];
    float vym = (y == 0) ? -drain : _Crest_VelocityYSource[uint2(x, y - 1)];

    // FIXME: Causes jitters which triggers foam.
    float hvel;
    if (_Crest_UpwindHeight)
    {
        // Get forward or backward diffs respectively.
        float h = _Crest_HeightSource[id];
        float2 hx = vxp <= 0.0
            ? float2(h, _Crest_HeightSource[id + uint2(1, 0)])
            : float2(_Crest_HeightSource[id - uint2(1, 0)], h);
        float2 hy = vyp <= 0.0
            ? float2(h, _Crest_HeightSource[id + uint2(0, 1)])
            : float2(_Crest_HeightSource[id - uint2(0, 1)], h);

        if (_Crest_DepthLimiter)
        {
            const float h_adj = HeightAdjustment(id);
            hx -= h_adj;
            hy -= h_adj;
        }

        hvel = -(vxp * hx[1] - vxm * hx[0] + vyp * hy[1] - vym * hy[0]) / _Crest_TexelSize;
    }
    else
    {
        const float dx = (vxp - vxm) / _Crest_TexelSize;
        const float dy = (vyp - vym) / _Crest_TexelSize;
        const float divergence = dx + dy;

        if (_Crest_DepthLimiter)
        {
            hvel = -(h - HeightAdjustment(id)) * divergence;
        }
        else
        {
            hvel = -h * divergence;
        }
    }

    h += _Crest_SimulationDeltaTime * hvel + pump - _Crest_Evaporation;
    h = max(h, 0.0);

    // // Dry - force to zero.
    // if (h < 0.001) h = 0.0;

    _Crest_HeightTarget[id] = h;
}


//
// Height Overshoot Reduction
//

void HeightOvershootReduction(const uint3 i_ID)
{
    const uint2 id = i_ID.xy;

    float h = _Crest_HeightSource[id];

    if (any(id == 0) || any(id == _Crest_Resolution - 1))
    {
        // Don't do this at boundary
        _Crest_HeightTarget[id] = h;
        return;
    }

    float y = _Crest_HeightSource[id] + _Crest_GroundHeightSource[id];
    float hn = _Crest_HeightSource[id + uint2(1, 0)];
    float hs = _Crest_HeightSource[id - uint2(1, 0)];
    float he = _Crest_HeightSource[id + uint2(1, 0)];
    float hw = _Crest_HeightSource[id - uint2(1, 0)];
    float yn = hn + _Crest_GroundHeightSource[id + uint2(1, 0)];
    float ys = hs + _Crest_GroundHeightSource[id - uint2(1, 0)];
    float ye = he + _Crest_GroundHeightSource[id + uint2(1, 0)];
    float yw = hw + _Crest_GroundHeightSource[id - uint2(1, 0)];

    float lambdaEdge = 2.0 * _Crest_TexelSize;

    // X axis
    if ((y - yw) > lambdaEdge && y > ye)
    {
        h += _Crest_OvershootReductionStrength * (max(0.0, 0.5 * (h + he)) - h);
    }
    if ((y - ye) > lambdaEdge && y > yw)
    {
        h += _Crest_OvershootReductionStrength * (max(0.0, 0.5 * (h + hw)) - h);
    }

    // Z axis
    if ((y - ys) > lambdaEdge && y > yn)
    {
        h += _Crest_OvershootReductionStrength * (max(0.0, 0.5 * (h + hn)) - h);
    }
    if ((y - yn) > lambdaEdge && y > ys)
    {
        h += _Crest_OvershootReductionStrength * (max(0.0, 0.5 * (h + hs)) - h);
    }

    _Crest_HeightTarget[id] = h;
}


//
// Update Velocity
//

void UpdateVelocity(const uint3 i_ID)
{
    const int2 id = i_ID.xy;

    // Clear velocity at border to prevent water leaking.
    if (i_ID.x == 0 || i_ID.y == 0 || i_ID.x == _Crest_Resolution - 1 || i_ID.y == _Crest_Resolution - 1)
    {
        _Crest_VelocityXTarget[id] = 0;
        _Crest_VelocityYTarget[id] = 0;
        return;
    }

    // Height before vel
    const float2 worldXZ_H0 = IdToWorldXZ_H(id);
    const float g0 = _Crest_GroundHeightSource[id];
    const float h0 = _Crest_HeightSource[id];

    const uint x = id.x;
    const uint y = id.y;
    const uint xp = min(x + 1, _Crest_Resolution - 1);
    const uint yp = min(y + 1, _Crest_Resolution - 1);

    const float gravity = 9.81f;

    // Vx
    {
        float vx = _Crest_VelocityXTarget[id];

        const float2 worldXZ_H1 = IdToWorldXZ_H(uint2(xp, y));
        const float g1 = _Crest_GroundHeightSource[uint2(xp, y)];
        const float h1 = _Crest_HeightSource[uint2(xp, y)];

        const float hdiff = (h1 + g1) - (h0 + g0);
        //if (fabs(hdiff) > maxDiff) hdiff = sign(hdiff) * maxDiff;

        float accel = -gravity * hdiff / _Crest_TexelSize;

        vx += _Crest_SimulationDeltaTime * accel;

        // Kill velocity if totally dry
        if (h0 < 0.001 && h1 < 0.001)
        {
            vx = 0.0;
        }
        else if (accel < 0 && h1 < 0.001)
        {
            // Don't produce acceleration from water flowing in from a dry cell
            vx = 0.0;
        }
        else if (accel > 0 && h0 < 0.001)
        {
            // Don't produce acceleration from water flowing in from a dry cell
            vx = 0.0;
        }

        _Crest_VelocityXTarget[id] = vx;

        if (h0 < 0.001 && h1 < 0.001) _Crest_VelocityXTarget[id] = 0.0;
    }

    // Vy
    {
        float vy = _Crest_VelocityYTarget[id];

        const float2 worldXZ_H1 = IdToWorldXZ_H(uint2(x, yp));
        const float g1 = _Crest_GroundHeightSource[uint2(x, yp)];
        const float h1 = _Crest_HeightSource[uint2(x, yp)];

        const float hdiff = (h1 + g1) - (h0 + g0);
        //if (fabs(hdiff) > maxDiff) hdiff = sign(hdiff) * maxDiff;

        float accel = -gravity * hdiff / _Crest_TexelSize;

        vy += _Crest_SimulationDeltaTime * accel;

        // Kill velocity if totally dry
        if (h0 < 0.001 && h1 < 0.001)
        {
            vy = 0.0;
        }
        else if (accel < 0 && h1 < 0.001)
        {
            // Don't produce acceleration from water flowing in from a dry cell
            vy = 0.0;
        }
        else if (accel > 0 && h0 < 0.001)
        {
            // Don't produce acceleration from water flowing in from a dry cell
            vy = 0.0;
        }

        _Crest_VelocityYTarget[id] = vy;
    }
}


//
// Blur
//

void Blur(const uint3 i_ID)
{
    const int2 id = i_ID.xy;
    const int resolution = _Crest_Resolution;

    float result = 0.0;
    const int radius = 1;
    float sampleCount = 0.0;
    for (int y = -radius; y <= radius; y++)
    {
        for (int x = -radius; x <= radius; x++)
        {
            int2 idx = id + int2(x, y);
            // Skip pixel if outside. Clamping coordinates produced bad values at edge.
            if (any(idx < 0) || any(idx >= resolution)) continue;
            result += _Crest_MaskSource[idx];
            sampleCount += 1.0;
        }
    }

    _Crest_MaskTarget[id] = result / sampleCount;
}

void BlurHeight(const uint3 i_ID)
{
    const int2 id = i_ID.xy;

    float result = 0.0;
    float twt = 0.0;
    const int rad = 1;

    float hcenter = _Crest_HeightSource[id];
    if (hcenter < 0.001)
    {
        _Crest_HeightTarget[id] = 0.0;
        return;
    }

    for (int y = -rad; y <= rad; y++)
    {
        for (int x = -rad; x <= rad; x++)
        {
            int2 idx;
            // Need to cast otherwise load will wrap around at edges.
            idx.x = clamp(id.x + x, 0, (int)_Crest_Resolution - 1);
            idx.y = clamp(id.y + y, 0, (int)_Crest_Resolution - 1);

            float h = _Crest_HeightSource[idx];
            if (h >= 0.001)
            {
                result += h + _Crest_GroundHeightSource[idx];
                twt += 1.0;
            }
        }
    }

    if (twt > 0.0)
    {
        float blurredY = result / twt;
        _Crest_HeightTarget[id] = max(blurredY - _Crest_GroundHeightSource[id], 0.0);
    }
    else
    {
        _Crest_HeightTarget[id] = hcenter;
    }
}

void AddAboveGroundHeight(const uint2 i_Coordinates, const float i_Ground, inout float io_Height, inout float io_Samples)
{
    float h = _Crest_HeightSource[i_Coordinates];

    if (h > FLT_MIN)
    {
        // Make sure we do not push water through ground.
        if (h < i_Ground)
        {
            io_Height += h;
            io_Samples += 1.0;
        }
    }
}

void MaskOut(const uint2 i_Coordinates, const float i_HeightBelowGround, inout bool i_Mask)
{
    const float h = _Crest_HeightSource[i_Coordinates];

    if (h > 0.0)
    {
        const float g = _Crest_GroundHeightSource[i_Coordinates];
        i_Mask = i_Mask || (i_HeightBelowGround < (h + g));
    }
}

void MaskEdge(const uint3 i_ID)
{
    const uint2 id = i_ID.xy;

    float h = _Crest_HeightSource[id];

    // Mask above ground height so we do not process it in expand edge.
    bool mask = h > 0.0;

    // Add ground height to water height to get world height of surface.
    h += _Crest_GroundHeightSource[id];

    if (!mask)
    {
        const int3 dd = int3(1, -1, 0);
        MaskOut(id + dd.xx, h, mask);
        MaskOut(id + dd.xy, h, mask);
        MaskOut(id + dd.xz, h, mask);
        MaskOut(id + dd.yx, h, mask);
        MaskOut(id + dd.yy, h, mask);
        MaskOut(id + dd.yz, h, mask);
        MaskOut(id + dd.zx, h, mask);
        MaskOut(id + dd.zy, h, mask);
    }

    if (mask)
    {
        // Move to world space.
        h += _Crest_DomainOrigin.y;
        // Make relative to sea level.
        h -= g_Crest_WaterCenter.y;
    }
    else
    {
        h = FLT_MIN;
    }

    _Crest_HeightTarget[id] = h;
}

void ExpandEdge(const uint3 i_ID)
{
    const uint2 id = i_ID.xy;
    const int3 dd = int3(1, -1, 0);

    float h = _Crest_HeightSource[id];

    if (h > FLT_MIN)
    {
        _Crest_HeightTarget[id] = h;
        return;
    }

    const float g = _Crest_GroundHeightSource[id] + _Crest_DomainOrigin.y - g_Crest_WaterCenter.y;

    // Add margin inside geometry to prevent rise or drops at intersections.

    float samples = 0.0;
    h = 0.0;

    AddAboveGroundHeight(id + dd.xx, g, h, samples);
    AddAboveGroundHeight(id + dd.xy, g, h, samples);
    AddAboveGroundHeight(id + dd.xz, g, h, samples);
    AddAboveGroundHeight(id + dd.yx, g, h, samples);
    AddAboveGroundHeight(id + dd.yy, g, h, samples);
    AddAboveGroundHeight(id + dd.yz, g, h, samples);
    AddAboveGroundHeight(id + dd.zx, g, h, samples);
    AddAboveGroundHeight(id + dd.zy, g, h, samples);

    // Found a height.
    if (samples > 0.0)
    {
        // Includes ground, sea level and origin.
        h /= samples;
    }

    _Crest_HeightTarget[id] = h;
}

void FinalizeHeight(const uint3 i_ID)
{
    const uint2 id = i_ID.xy;
    float h = _Crest_HeightSource[id];

    // Clip doesn't work if ground is below sea level - it will make the water pop back up to sea level/anim waves.
    // Same for multiplying down alpha blend weight. This does something reasonable - yanks down water when almost dry..
    h -= 0.1 * (1.0 - saturate(h / 0.02));

    // Add ground height to water height to get world height of surface.
    h += _Crest_GroundHeightSource[id];

    // Move to world space.
    h += _Crest_DomainOrigin.y;
    // Make relative to sea level.
    h -= g_Crest_WaterCenter.y;

    _Crest_HeightTarget[id] = h;
}


//
// Injectors
//

float WaterHeight(float2 i_UV, float i_Current = 0.0)
{
    float h = _Crest_HeightSource.SampleLevel(LODData_linear_clamp_sampler, i_UV, 0.0).x;

    // Reject any samples below ground.
    if (h < 0.001)
    {
        // Return the current height to not trigger the curve test.
        h = i_Current;
    }
    else
    {
        // Add ground height to water height to get height of surface
        h += _Crest_GroundHeightSource.SampleLevel(LODData_linear_clamp_sampler, i_UV, 0.0).x;
    }

    // Silences: shader warning use of potentially uninitialized variable.
    return h;
}

float ComputeAlpha(const float2 i_UV)
{
    // Domain mask.
    const float2 offset = abs(i_UV - 0.5);
    float alpha = smoothstep(0.5, 0.45, max(offset.x, offset.y));

    if (alpha >= 0.0)
    {
        // Simulation mask.
        alpha *= _Crest_MaskSource.SampleLevel(LODData_linear_clamp_sampler, i_UV, 0.0);
    }

    return alpha;
}

float ComputeAlphaFromMargin(const float2 i_Position)
{
    float alpha = 1.0;
    float2 position = abs(i_Position - _Crest_DomainOrigin.xz);
    if (max(position.x, position.y) > (_Crest_DomainWidth.x * 0.5 - _Crest_MarginWidth)) alpha = 0.0;
    return alpha;
}

void InjectShape(const uint3 i_ID)
{
    const float2 position = Cascade::MakeAnimatedWaves(i_ID.z).IDToWorld(i_ID.xy);
    const float2 uv = (position - _Crest_DomainOrigin.xz) / _Crest_DomainWidth + 0.5;

    // Over scan to ensure signal continued off the edges which helps at low LODs.
    if (any(uv != saturate(uv))) return;

    float alpha = _Crest_Weight * ComputeAlpha(uv);
    if (alpha <= 0.0) return;

    float h = _Crest_HeightSource.SampleLevel(LODData_linear_clamp_sampler, uv, 0.0).x;

    // Power up alpha to bring anim waves further in towards shore.
    alpha = pow(alpha, 2.0);

    float4 result = _Crest_ShapeTarget[i_ID];

    // Alpha blend.
    result.xyz *= 1.0 - alpha;
    result.xyz += float3(0, h * alpha, 0);

    _Crest_ShapeTarget[i_ID] = result;
}

void InjectLevel(const uint3 i_ID)
{
    const float2 position = Cascade::MakeLevel(i_ID.z).IDToWorld(i_ID.xy);
    const float2 uv = (position - _Crest_DomainOrigin.xz) / _Crest_DomainWidth + 0.5;

    // Over scan to ensure signal continued off the edges which helps at low LODs.
    if (any(uv != saturate(uv))) return;

    float alpha = _Crest_Weight * ComputeAlphaFromMargin(position);
    if (alpha <= 0.0) return;

    float h = _Crest_HeightSource.SampleLevel(LODData_linear_clamp_sampler, uv, 0.0).x;

    float result = _Crest_Weight * _Crest_LevelTarget[i_ID];

    // Height is a source of water and should be the minimum height.
    // The simulation will still ripple over the top which is acceptable.
    result = max(result, h);

    _Crest_LevelTarget[i_ID] = result;
}

void InjectFoam(const uint3 i_ID)
{
    const float2 position = Cascade::MakeFoam(i_ID.z).IDToWorld(i_ID.xy);
    const float2 uv = (position - _Crest_DomainOrigin.xz) / _Crest_DomainWidth + 0.5;

    // Over scan to ensure signal continued off the edges which helps at low LODs.
    if (any(uv != saturate(uv))) return;

    float alpha = _Crest_Weight * (_Crest_InjectLevel ? ComputeAlphaFromMargin(position) : ComputeAlpha(uv));
    if (alpha <= 0.0) return;

    // Use approximation of max curvature as foam term. Seems to grab leading wave edge alright.
    const float2 dx = float2(1.0 / _Crest_Resolution, 0.0);
    const float h = WaterHeight(uv);
    if (h <= 0.0) return;
    const float h_xm = WaterHeight(uv - dx.xy, h);
    const float h_xp = WaterHeight(uv + dx.xy, h);
    const float h_zm = WaterHeight(uv - dx.yx, h);
    const float h_zp = WaterHeight(uv + dx.yx, h);
    const float curvature = max(abs(h_xp + h_xm - 2.0 * h), abs(h_zp + h_zm - 2.0 * h)) / (2.0 * dx.x * _Crest_DomainWidth);

    float foam = 10.0 * curvature;

    // Apply velocity to flow to prevent foam at edges of geometry.
    const float vx = _Crest_VelocityXSource.SampleLevel(LODData_linear_clamp_sampler, uv, 0.0);
    const float vy = _Crest_VelocityYSource.SampleLevel(LODData_linear_clamp_sampler, uv, 0.0);
    foam *= length(float2(vx, vy));

    foam *= alpha;

    // Integrate.
    foam *= _Crest_SimDeltaTime;

    foam *= _Crest_InjectionStrength;

    _Crest_FoamTarget[i_ID] += foam;
}

void InjectFlow(const uint3 i_ID)
{
    const Cascade cascade = Cascade::MakeFlow(i_ID.z);
    const float2 position = cascade.IDToWorld(i_ID.xy);
    const float2 uv = (position - _Crest_DomainOrigin.xz) / _Crest_DomainWidth + 0.5;

    // Over scan to ensure signal continued off the edges which helps at low LODs.
    if (any(uv != saturate(uv))) return;

    float alpha = _Crest_Weight * (_Crest_InjectLevel ? ComputeAlphaFromMargin(position) : ComputeAlpha(uv));
    if (alpha <= 0.0) return;

    float mip;
    {
        const float3 uv = cascade.IDToUV(i_ID.xy);
        const float2 dd = float2(cascade._OneOverResolution * 8.0, 0);
        const float2 dx = cascade.UVToWorld(uv + dd.xyy) - position;
        const float2 dy = cascade.UVToWorld(uv + dd.yxy) - position;
        // Calculate the rate of change in x and y directions.
        const float2 rateOfChange = float2(length(dx.xy), length(dy.xy));
        mip = log2(max(rateOfChange.x, rateOfChange.y));
    }

    // These should use mip maps. Without this bad pops ensue from aliasing x combine-pass-flow.
    const float vx = _Crest_VelocityXSource.SampleLevel(LODData_linear_clamp_sampler, uv, mip);
    const float vy = _Crest_VelocityYSource.SampleLevel(LODData_linear_clamp_sampler, uv, mip);

    float2 result = _Crest_FlowTarget[i_ID];

    // Alpha blend.
    result *= 1.0 - alpha;
    result += float2(vx, vy) * _Crest_InjectionStrength * alpha;

    _Crest_FlowTarget[i_ID] = result;
}


//
// Baking
//

// Copy texture. Here because script code path simpler if same for all outputs
// rather than using CopyTexture.
void BakeLevel(const uint3 i_ID)
{
    const uint2 id = i_ID.xy;
    _Crest_LevelBakeTarget[id] = _Crest_HeightSource[id];
}

void BakeFoam(const uint3 i_ID)
{
    const uint2 id = i_ID.xy;
    const float2 uv = id / (float)_Crest_Resolution;

    // Over scan to ensure signal continued off the edges which helps at low LODs.
    if (any(uv != saturate(uv))) return;

    float alpha = ComputeAlphaFromMargin(IdToWorldXZ_H(id));
    if (alpha <= 0.0) return;

    // Cannot load texture as curvature is not picked up. Must need interpolation.
    const float2 dx = float2(1.0 / _Crest_Resolution, 0.0);
    const float h = WaterHeight(uv);
    if (h <= 0.0)
    {
        _Crest_FoamBakeTarget[id] = 0.0;
        return;
    }
    const float h_xm = WaterHeight(uv - dx.xy, h);
    const float h_xp = WaterHeight(uv + dx.xy, h);
    const float h_zm = WaterHeight(uv - dx.yx, h);
    const float h_zp = WaterHeight(uv + dx.yx, h);
    const float curvature = max(abs(h_xp + h_xm - 2.0 * h), abs(h_zp + h_zm - 2.0 * h)) / (2.0 * dx.x * _Crest_DomainWidth);

    float foam = 10.0 * curvature;

    foam *= length(float2(_Crest_VelocityXSource[id], _Crest_VelocityYSource[id]));

    foam *= _Crest_InjectionStrength;

    _Crest_FoamBakeTarget[id] = foam;
}

void BakeFlow(const uint3 i_ID)
{
    const uint2 id = i_ID.xy;
    float alpha = ComputeAlphaFromMargin(IdToWorldXZ_H(id));
    if (alpha <= 0.0) return;
    _Crest_FlowBakeTarget[id.xy] = float2(_Crest_VelocityXSource[id], _Crest_VelocityYSource[id]) * _Crest_InjectionStrength;
}

m_CrestNameSpaceEnd

m_CrestKernelDefault(Initialize);
m_CrestKernelDefault(InitializeGroundHeight);
m_CrestKernelDefault(Advect);
m_CrestKernelDefault(UpdateHeight);
m_CrestKernelDefault(HeightOvershootReduction);
m_CrestKernelDefault(UpdateVelocity);
m_CrestKernelDefault(BlurHeight);
m_CrestKernelDefault(Blur);
m_CrestKernelDefault(MaskEdge);
m_CrestKernelDefault(ExpandEdge);
m_CrestKernelDefault(FinalizeHeight);

m_CrestInputKernelDefault(InjectShape);
m_CrestInputKernelDefault(InjectLevel);
m_CrestInputKernelDefault(InjectFoam);
m_CrestInputKernelDefault(InjectFlow);

m_CrestKernelDefault(BakeLevel);
m_CrestKernelDefault(BakeFoam);
m_CrestKernelDefault(BakeFlow);