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

// Inspired by https://github.com/speps/GX-EncinoWaves

#pragma exclude_renderers glcore gles3

#pragma kernel SpectrumInitalize
#pragma kernel SpectrumUpdate

#pragma multi_compile_local _ d_AdvancedControls

#define INV2PI	0.15915494309f
#define PI4		0.33661977236f
#define INVPI2	0.63661977236f
#define HPI		1.57079632679f
#define PI		3.14159265358f
#define PI2		6.28318530717f
#define HSQRT2	0.70710678118f

// These must match corresponding constants in WaveSpectrum.cs
#define SPECTRUM_OCTAVE_COUNT_INT	14
#define SPECTRUM_OCTAVE_COUNT		14.0
#define SPECTRUM_SMALLEST_WL_POW_2	-4.0

#include "Packages/com.waveharmonic.crest/Runtime/Shaders/Library/Macros.hlsl"

uint _Crest_Size;
float _Crest_WindSpeed;
float _Crest_Turbulence;
float _Crest_Gravity;
float _Crest_Period;
float _Crest_Alignment;

uint WangHash(uint seed)
{
    seed = (seed ^ 61) ^ (seed >> 16);
    seed *= 9;
    seed = seed ^ (seed >> 4);
    seed *= 0x27d4eb2d;
    seed = seed ^ (seed >> 15);
    return seed;
}

uint Rand(inout uint rngState)
{
    rngState ^= (rngState << 13);
    rngState ^= (rngState >> 17);
    rngState ^= (rngState << 5);
    return rngState;
}

float RandFloat(inout uint rngState)
{
    return Rand(rngState) / 4294967296.0f;
}

float RandGauss(inout uint rngState)
{
    float u1 = RandFloat(rngState);
    float u2 = RandFloat(rngState);
    if (u1 < 1e-6f)
        u1 = 1e-6f;
    return sqrt(-2.0f * log(u1)) * cos(PI2 * u2);
}

void DeepDispersion(float k, out float w, out float dwdk)
{
    w = sqrt(abs(_Crest_Gravity * k));

    // Allow FFT to loop in time
    if( _Crest_Period > 0.0 )
    {
        float thisPeriod = PI2 / w;
        float loops = _Crest_Period / thisPeriod;
        // Make sure loops integral number of times
        loops = ceil( loops );
        // Work our way back to frequency
        thisPeriod = _Crest_Period / loops;
        w = PI2 / thisPeriod;
    }

    dwdk = _Crest_Gravity / (2.0f * w);
}

float AlphaBetaSpectrum(float A, float B, float g, float w, float wm)
{
    return
        (A * g * g / pow(w, 5.0f)) *
        exp(-B * pow(wm / w, 4.0f));
}

float PiersonMoskowitzSpectrum(float w)
{
    float wm = 0.87f * _Crest_Gravity / _Crest_WindSpeed;
    return AlphaBetaSpectrum(8.1e-3f, 1.291f, _Crest_Gravity, w, wm);
}

float PiersonMoskowitzWindTerm( float w )
{
    float wm = 0.87f * _Crest_Gravity / _Crest_WindSpeed;
    return exp( -1.291 * pow( wm / w, 4.0f ) );
}

float PosCosSquaredDirectionalSpreading( float cosTheta )
{
    // Aligned waves.
    float alignment;
    {
        float minWeight = lerp(1.0, 0.1, _Crest_Alignment);
        float wt = max(cosTheta, minWeight);
        wt *= lerp(0.25, 1.5, _Crest_Alignment);
        float power = lerp(1.0, 16.0, _Crest_Alignment);

        // Needs 2 to match current settings.
        alignment = wt * INVPI2 * pow(abs(cosTheta), power) * PI;
    }

    float turbulence;
    {
        if (cosTheta > 0.0)
        {
            turbulence = lerp(INVPI2 * (cosTheta * cosTheta), PI4, _Crest_Turbulence);
        }
        else
        {
            turbulence = PI4 * _Crest_Turbulence;
        }
    }

    // Integrate alignment and turbulence.
    return lerp(turbulence, alignment, _Crest_Alignment * (1.0 - _Crest_Turbulence));
}

RWTexture2DArray<float4> _Crest_ResultInit;
Texture2D<float> _Crest_SpectrumControls;
SamplerState linear_clamp_sampler;
float2 _Crest_WindDir;

[numthreads(8,8,1)]
void SpectrumInitalize(uint3 id : SV_DispatchThreadID)
{
    const int2 center = _Crest_Size.xx / 2;
    const int2 coord = id.xy - center;

    if( coord.x == 0 && coord.y == 0 )
    {
        _Crest_ResultInit[id] = 0.0;
        return;
    }

    uint depth;
    {
        uint width, height;
        _Crest_ResultInit.GetDimensions( width, height, depth );
    }

    uint maxCoord = int( max( abs( coord.x ), abs( coord.y ) ) );

    // Matches variable with same name in ShapeFFT.cs
    uint WAVE_SAMPLE_FACTOR = 8;

    // If not largest cascade which will get all wavelengths, then limit
    // so we split range of frequencies with no overlaps.
    // The check on maxCoord below looks pretty magic. It is optimised version of:
    // uint samplesPerWave = _Crest_Size / WAVE_SAMPLE_FACTOR;
    // Too low wavelength (maxCoord < _Crest_Size / (2 * samplesPerWave) ||
    // Too high wavelength maxCoord >= _Crest_Size / samplesPerWave)
    if ( id.z < (depth - 1) &&
        // Too low wavelength
        maxCoord < WAVE_SAMPLE_FACTOR / 2 ||
        // Too high wavelength
        maxCoord >= WAVE_SAMPLE_FACTOR
        )
    {
        _Crest_ResultInit[id] = 0.0;
        return;
    }

    const float worldSize = 0.5f * (1 << id.z);

    // Find wave vector and number
    const float2 k = PI2 * coord / worldSize;
    const float kMag = length(k);

    // Init seed. rngState was _RngState (file level variable), but GameCore platforms does not allow assigning to them.
    // See issue #856 for more information: https://github.com/wave-harmonic/crest/issues/856
    uint rngState = WangHash(id.z * _Crest_Size * _Crest_Size + id.y * _Crest_Size + id.x);

    // Dispersion
    float w; float dwdk;
    DeepDispersion(kMag, w, dwdk);

    // Spectrum - use power values from users spectrum, but borrow wind term from PM
    const float wavelength = PI2 / kMag;
    const float octaveIndex = log2( wavelength ) - SPECTRUM_SMALLEST_WL_POW_2;
    const float2 spectrumUV = float2((octaveIndex + 0.5) / SPECTRUM_OCTAVE_COUNT, 0.5);
    const float spectrum = _Crest_SpectrumControls.SampleLevel( linear_clamp_sampler, spectrumUV, 0.0 ) *
        PiersonMoskowitzWindTerm( w );

    float deltaSPos = spectrum;
    float deltaSNeg = spectrum;

    // Directional spreading
    const float cosTheta = dot( k, _Crest_WindDir ) / kMag;
    deltaSPos *= PosCosSquaredDirectionalSpreading( cosTheta );
    deltaSNeg *= PosCosSquaredDirectionalSpreading( -cosTheta );
    const float dK = PI2 / worldSize;
    deltaSPos *= (dK * dK) * dwdk / kMag;
    deltaSNeg *= (dK * dK) * dwdk / kMag;

    // Amplitude
    const float ampPos = RandGauss(rngState) * sqrt(abs(deltaSPos) * 2.0f);
    const float ampNeg = RandGauss(rngState) * sqrt(abs(deltaSNeg) * 2.0f);

    // Output
    const float phasePos = RandFloat(rngState) * PI2;
    const float phaseNeg = RandFloat(rngState) * PI2;

    //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    const float spiceyMultiplier = 1.5;
    _Crest_ResultInit[id] = float4(ampPos * float2(cos(phasePos), -sin(phasePos)), ampNeg * float2(cos(phaseNeg), -sin(phaseNeg)))
        * spiceyMultiplier;
}

float _Crest_Time;
float _Crest_Chop;

#if d_AdvancedControls
float _Crest_ChopScales[SPECTRUM_OCTAVE_COUNT_INT];
float _Crest_GravityScales[SPECTRUM_OCTAVE_COUNT_INT];
#endif

Texture2DArray<float4> _Crest_Init0;
RWTexture2DArray<m_Float2> _Crest_ResultHeight;
RWTexture2DArray<m_Float2> _Crest_ResultDisplaceX;
RWTexture2DArray<m_Float2> _Crest_ResultDisplaceZ;

float2 cmul(float2 lhs, float2 rhs)
{
    return float2(
        lhs.x * rhs.x - lhs.y * rhs.y,
        lhs.x * rhs.y + lhs.y * rhs.x
    );
}

[numthreads(8, 8, 1)]
void SpectrumUpdate(uint3 id : SV_DispatchThreadID)
{
    const int2 center = _Crest_Size.xx / 2;
    const int2 coord = id.xy - center;

    // Find wave vector and number
    const float worldSize = 0.5 * (1 << id.z);
    const float2 k = PI2 * coord / worldSize;
    const float kMag = length(k);

    float time = _Crest_Time;
    float chop = _Crest_Chop;

#if d_AdvancedControls
    const float wavelength = PI2 / kMag;
    const uint octaveIndex = clamp(log2(wavelength) - SPECTRUM_SMALLEST_WL_POW_2, 0.0, SPECTRUM_OCTAVE_COUNT - 1.0);
    time *= _Crest_GravityScales[octaveIndex];
    chop *= _Crest_ChopScales[octaveIndex];
#endif

    // Dispersion
    float w; float dwdk;
    DeepDispersion(kMag, w, dwdk);

    // Advance time
    float sw; float cw;
    sincos(w * time, sw, cw);

    const float2 fwd = float2(cw, -sw);
    const float2 bkwd = float2(cw, sw);

    const float4 h0 = _Crest_Init0[id];
    const float2 h = cmul(h0.xy, fwd) + cmul(h0.zw, bkwd);

    _Crest_ResultHeight[id]    = m_Float2FromFloat2(h);
    _Crest_ResultDisplaceX[id] = m_Float2FromFloat2((chop * float2(-h.y * k.x, h.x * k.x) / (kMag + 0.00001f)));
    _Crest_ResultDisplaceZ[id] = m_Float2FromFloat2((chop * float2(-h.y * k.y, h.x * k.y) / (kMag + 0.00001f)));
}