import lib.util;

[[vk::binding(0, 0)]]
StructuredBuffer<float2> input;

[[vk::binding(1, 0)]]
RWStructuredBuffer<float2> output;

[[vk::binding(2, 0)]]
StructuredBuffer<float2> twiddles;

[[vk::binding(3, 0)]]
cbuffer Uniforms {
  // Invariant: stride * big_n = input_len
  uint stride;
  uint big_n;
  uint input_len;
};

// 256 threads per FFT
[shader("compute")]
[numthreads(256, 1, 1)]
void fft_v1_batched(uint3 globalId : SV_DispatchThreadID)
{
    const uint id = globalId.x;
    const uint stride = Uniforms.stride;
    const uint big_n = Uniforms.big_n;
    const uint half_big_n = big_n / 2;
    const uint k = id / stride;

    const uint which_fft = globalId.y;
    const uint start_idx = which_fft * Uniforms.input_len + id % stride;

    // Get odd and even elements.
    let p = input[start_idx + 2 * k * stride];
    let q = input[start_idx + (2 * k + 1) * stride];
    // Combine.
    let factor = twiddles[k * stride];
    let qf = cmul(q, factor);
    output[start_idx + k * stride] = p + qf;
    output[start_idx + (k + half_big_n) * stride] = p - qf;
}