import lib.util;

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

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

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

[[vk::binding(3, 0)]]
cbuffer Uniforms {
  uint input_len;
  uint n_iter;
};

void fft_v2_stage(uint which_fft, uint id, uint stride, uint big_n, RWStructuredBuffer<float2> in, RWStructuredBuffer<float2> out) {
    const uint k = id / stride;
    const uint start_idx = which_fft * Uniforms.input_len + id % stride;

    // Collect inputs.
    let i0 = in[start_idx + 4 * k * stride];
    let i1 = in[start_idx + (4 * k + 1) * stride];
    let i2 = in[start_idx + (4 * k + 2) * stride];
    let i3 = in[start_idx + (4 * k + 3) * stride];
    // Collect relevant twiddles.
    let ot1 = twiddles[1 * k * stride];
    let ot2 = twiddles[2 * k * stride];
    let ot3 = twiddles[3 * k * stride];

    let a = i0;
    let b = cmul(ot1, i1);
    let c = cmul(ot2, i2);
    let d = cmul(ot3, i3);

    // To derive this, write the output assignments in terms of
    // a/b/c/d, then factor out!
    let ac_sum = a + c;
    let ac_diff = a - c;
    let bd_sum = b + d;
    let bd_diff_ni = mul_ni(b - d);

    out[start_idx + k * stride] = ac_sum + bd_sum;
    out[start_idx + (k + big_n / 4) * stride] = ac_diff + bd_diff_ni;
    out[start_idx + (k + big_n / 2) * stride] = ac_sum - bd_sum;
    out[start_idx + (k + 3 * big_n / 4) * stride] = ac_diff - bd_diff_ni;
}

[shader("compute")]
[numthreads(64, 16, 1)]
void fft_v2_barrier(uint3 globalId : SV_DispatchThreadID, uint id: SV_GroupIndex)
{
    const uint which_fft = globalId.z;

    uint stride = Uniforms.input_len;
    uint big_n = 1;

    for (uint i = 0; i < Uniforms.n_iter; i++) {
      stride /= 4;
      big_n *= 4;
      if (i % 2 == 0) {
        fft_v2_stage(which_fft, id, stride, big_n, input, output);
      } else {
        fft_v2_stage(which_fft, id, stride, big_n, output, input);
      }
      AllMemoryBarrierWithGroupSync();
    }
}