#version 310 es

layout(local_size_x = 64) in;

layout(std430, binding = 0) readonly buffer Distribution
{
    vec2 distribution[];
};

layout(std430, binding = 1) writeonly buffer HeightmapFFT
{
    uint heights[];
};

layout(binding = 2, std140) uniform UBO
{
    vec4 uModTime;
};

vec2 alias(vec2 i, vec2 N)
{
    return mix(i, i - N, greaterThan(i, 0.5 * N));
}

vec4 cmul(vec4 a, vec4 b)
{
    vec4 r3 = a.yxwz;
    vec4 r1 = b.xxzz;
    vec4 R0 = a * r1;
    vec4 r2 = b.yyww;
    vec4 R1 = r2 * r3;
    return R0 + vec4(-R1.x, R1.y, -R1.z, R1.w);
}

vec2 cmul(vec2 a, vec2 b)
{
    vec2 r3 = a.yx;
    vec2 r1 = b.xx;
    vec2 R0 = a * r1;
    vec2 r2 = b.yy;
    vec2 R1 = r2 * r3;
    return R0 + vec2(-R1.x, R1.y);
}

uint pack2(vec2 v)
{
    return packHalf2x16(v);
}

uvec2 pack4(vec4 v)
{
    return uvec2(packHalf2x16(v.xy), packHalf2x16(v.zw));
}

uvec2 workaround_mix(uvec2 a, uvec2 b, bvec2 sel)
{
    return uvec2(sel.x ? b.x : a.x, sel.y ? b.y : a.y);
}

void generate_heightmap()
{
    uvec2 N = gl_WorkGroupSize.xy * gl_NumWorkGroups.xy;
    uvec2 i = gl_GlobalInvocationID.xy;
    // Pick out the negative frequency variant.
    uvec2 wi = workaround_mix(N - i, uvec2(0u), equal(i, uvec2(0u)));

    // Pick out positive and negative travelling waves.
    vec2 a = distribution[i.y * N.x + i.x];
    vec2 b = distribution[wi.y * N.x + wi.x];

    vec2 k = uModTime.xy * alias(vec2(i), vec2(N));
    float k_len = length(k);

    const float G = 9.81;

    // If this sample runs for hours on end, the cosines of very large numbers will eventually become unstable.
    // It is fairly easy to fix this by wrapping uTime,
    // and quantizing w such that wrapping uTime does not change the result.
    // See Tessendorf's paper for how to do it.
    // The sqrt(G * k_len) factor represents how fast ocean waves at different frequencies propagate.
    float w = sqrt(G * k_len) * uModTime.z;
    float cw = cos(w);
    float sw = sin(w);

    // Complex multiply to rotate our frequency samples.
    a = cmul(a, vec2(cw, sw));
    b = cmul(b, vec2(cw, sw));
    b = vec2(b.x, -b.y); // Complex conjugate since we picked a frequency with the opposite direction.
    vec2 res = a + b; // Sum up forward and backwards travelling waves.
    heights[i.y * N.x + i.x] = pack2(res);
}

void main()
{
    generate_heightmap();
}