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

// Compute shader to perform combine of displacements. Reads and writes to texture array which saves
// needing to do ping pong of render targets. Unfortunately reading/writing float4s is not supported
// on pre-DX11.3 hardware (aka typed UAV loads), so this path is not the default, for now..

#pragma kernel ShapeCombine
#pragma kernel ShapeCombine_DISABLE_COMBINE _DISABLE_COMBINE
#pragma kernel ShapeCombine_FLOW_ON _FLOW_ON
#pragma kernel ShapeCombine_FLOW_ON_DISABLE_COMBINE _FLOW_ON _DISABLE_COMBINE
#pragma kernel ShapeCombine_DYNAMIC_WAVE_SIM_ON _DYNAMIC_WAVE_SIM_ON
#pragma kernel ShapeCombine_DYNAMIC_WAVE_SIM_ON_DISABLE_COMBINE _DYNAMIC_WAVE_SIM_ON _DISABLE_COMBINE
#pragma kernel ShapeCombine_FLOW_ON_DYNAMIC_WAVE_SIM_ON _FLOW_ON _DYNAMIC_WAVE_SIM_ON
#pragma kernel ShapeCombine_FLOW_ON_DYNAMIC_WAVE_SIM_ON_DISABLE_COMBINE _FLOW_ON _DYNAMIC_WAVE_SIM_ON _DISABLE_COMBINE
#pragma kernel ShapeCombineDynamicWaves
#pragma kernel ShapeCombineDynamicWaves_DISABLE_COMBINE _DISABLE_COMBINE
#pragma kernel ShapeCombineCopyDynamicWaves

#include "HLSLSupport.cginc"

#include "Packages/com.waveharmonic.crest/Runtime/Shaders/Library/Macros.hlsl"
#include "Packages/com.waveharmonic.crest/Runtime/Shaders/Library/Constants.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/Cascade.hlsl"

float _Crest_HorizontalDisplace;
float _Crest_DisplaceClamp;

Texture2DArray _Crest_WaveBuffer;
RWTexture2DArray<float4> _Crest_Target;
RWTexture2DArray<float4> _Crest_DynamicWavesTarget;
RWTexture2DArray<float4> _Crest_AnimatedWavesTarget;

m_CrestNameSpace

void Flow(const float texel, out float2 offsets, out float2 weights)
{
    const float period = max(3.0 * texel, 1.0);
    const float half_period = period / 2.0;
    offsets = fmod(float2(g_Crest_Time, g_Crest_Time + half_period), period);
    weights.x = offsets.x / half_period;
    if (weights.x > 1.0) weights.x = 2.0 - weights.x;
    weights.y = 1.0 - weights.x;
}

void SampleDisplacementsCompute(
    in RWTexture2DArray<float4> i_dispSampler,
    in float i_resolution, in float3 i_uv_slice,
    in float i_wt, inout float3 io_worldPos
) {
    // NOTE: We have to roll our own bilinear filter in Compute shaders when
    // reading from a RWTexture. The documentation below explains how SRV
    // and UAV mappings of the same texture cannot exist at the same time.
    // https://docs.microsoft.com/en-us/windows/desktop/direct3dhlsl/sm5-object-rwtexture2d

    // Convert from UV to coordinates
    const float2 pixelCoord = i_uv_slice.xy * i_resolution;

    // Make relative to pixel centers
    float2 pixelCoordCenters = pixelCoord - 0.5;

    // Clamp from below and above (desired?)
    pixelCoordCenters = clamp(pixelCoordCenters, 0.0, i_resolution - 1.0);

    // Compute integral and fractional parts
    const uint2 pixelCoordCentersBotLeft = floor(pixelCoordCenters);
    const uint sliceIndex = i_uv_slice.z;
    const float2 pixelCoordCentersFrac = frac(pixelCoordCenters);

    const half4 dataBotLeft = i_dispSampler[uint3(pixelCoordCentersBotLeft, sliceIndex)];
    const half4 dataBotRight = i_dispSampler[uint3(pixelCoordCentersBotLeft + uint2(1, 0), sliceIndex)];
    const half4 dataTopLeft = i_dispSampler[uint3(pixelCoordCentersBotLeft + uint2(0, 1), sliceIndex)];
    const half4 dataTopRight = i_dispSampler[uint3(pixelCoordCentersBotLeft + uint2(1, 1), sliceIndex)];

    const float4 dataLerped = lerp(
        lerp(dataBotLeft, dataBotRight, pixelCoordCentersFrac.x),
        lerp(dataTopLeft, dataTopRight, pixelCoordCentersFrac.x),
        pixelCoordCentersFrac.y
    );

    io_worldPos += i_wt * dataLerped.xyz;
}

void ShapeCombineBase(uint3 id)
{
    const uint slice0 = _Crest_LodIndex;
    const Cascade cascade = Cascade::MakeAnimatedWaves(slice0);

    const float3 uv = cascade.IDToUV(id.xy);
    const float2 positionWSXZ = cascade.UVToWorld(uv);

    float3 result = 0.0;

    // Sample in waves for this cascade.
    {
#if _FLOW_ON
        const half2 flow = Cascade::MakeFlow(slice0).SampleFlow(positionWSXZ);

        float2 offsets, weights;
        Flow(cascade._Texel, offsets, weights);

        result += cascade.Sample(_Crest_WaveBuffer, positionWSXZ - offsets[0] * flow).xyz * weights[0];
        result += cascade.Sample(_Crest_WaveBuffer, positionWSXZ - offsets[1] * flow).xyz * weights[1];
#else
        result += cascade.Sample(_Crest_WaveBuffer, uv).xyz;
#endif // _FLOW_ON
    }

    // Disabled for last LOD.
#if !_DISABLE_COMBINE
    {
        const Cascade cascade = Cascade::MakeAnimatedWaves(slice0 + 1);
        // Sample the shape 1 texture at this world position.
        const float3 uv = cascade.WorldToUV(positionWSXZ);
        // Waves to combine down from the next lod up the chain.
        SampleDisplacementsCompute(_Crest_Target, cascade._Resolution, uv, 1.0, result);
    }
#endif

#if _DYNAMIC_WAVE_SIM_ON
    {
        // Convert dynamic wave sim to displacements.
        result += Cascade::MakeDynamicWaves(slice0)
            .SampleDynamicWavesDisplacement(positionWSXZ, _Crest_HorizontalDisplace, _Crest_DisplaceClamp);
    }
#endif // _DYNAMIC_WAVE_SIM_ON

    _Crest_Target[uint3(id.xy, slice0)] = float4(result, 0.0);
}

void ShapeCombineDynamicWaves(uint3 id)
{
    const uint slice0 = _Crest_LodIndex;
    const Cascade cascade = Cascade::MakeAnimatedWaves(slice0);
    const float3 uv = cascade.IDToUV(id.xy);
    const float2 positionWSXZ = cascade.UVToWorld(uv);
    float3 result = 0.0;

    // Disabled for last LOD.
#if !_DISABLE_COMBINE
    {
        // We are sampling from the target which matches Animated Waves descriptor.
        const Cascade cascade = Cascade::MakeAnimatedWaves(slice0 + 1);
        const float3 uv = cascade.WorldToUV(positionWSXZ);
        // Waves to combine down from the next lod up the chain.
        SampleDisplacementsCompute(_Crest_DynamicWavesTarget, cascade._Resolution, uv, 1.0, result);
    }
#endif

    {
        // We are sampling from Dynamic Waves.
        const Cascade cascade = Cascade::MakeDynamicWaves(slice0);
        const float3 uv = cascade.WorldToUV(positionWSXZ);
        result += cascade.SampleDynamicWavesDisplacement(uv, _Crest_HorizontalDisplace, _Crest_DisplaceClamp);
    }

    _Crest_DynamicWavesTarget[uint3(id.xy, slice0)] = float4(result, 0.0);
}

void ShapeCombineCopyDynamicWaves(uint3 id)
{
    _Crest_AnimatedWavesTarget[id] += _Crest_DynamicWavesTarget[id];
}

m_CrestNameSpaceEnd

[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombine(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineBase(id); }
[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombine_DISABLE_COMBINE(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineBase(id); }
[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombine_FLOW_ON(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineBase(id); }
[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombine_FLOW_ON_DISABLE_COMBINE(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineBase(id); }
[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombine_DYNAMIC_WAVE_SIM_ON(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineBase(id); }
[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombine_DYNAMIC_WAVE_SIM_ON_DISABLE_COMBINE(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineBase(id); }
[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombine_FLOW_ON_DYNAMIC_WAVE_SIM_ON(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineBase(id); }
[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombine_FLOW_ON_DYNAMIC_WAVE_SIM_ON_DISABLE_COMBINE(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineBase(id); }
[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombineDynamicWaves(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineDynamicWaves(id); }
[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombineDynamicWaves_DISABLE_COMBINE(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineDynamicWaves(id); }
[numthreads(THREAD_GROUP_SIZE_X, THREAD_GROUP_SIZE_Y, 1)] void ShapeCombineCopyDynamicWaves(uint3 id : SV_DispatchThreadID) { m_Crest::ShapeCombineCopyDynamicWaves(id); }