#include <metal_stdlib>
using namespace metal;

/// Radix-2 Cooley-Tukey FFT with shared memory butterfly operations.
/// Supports both forward and inverse transforms.

struct FftParams {
    uint length;       // Length of the current 1D FFT
    uint log2_length;  // log2 of length
    uint stage;        // Current butterfly stage (0-based)
    uint inverse;      // 0 = forward, 1 = inverse
    uint total_count;  // Total number of elements (width * height for 2D)
    uint stride;       // Stride for 2D row/column decomposition
    uint is_column;    // 0 = row pass, 1 = column pass
    uint width;        // Image width (for column indexing)
};

// Butterfly operation on a pair of complex numbers
// x = a + W*b,  y = a - W*b
// where W = exp(-2*pi*i*k/N) for forward, exp(+2*pi*i*k/N) for inverse.

kernel void fft_radix2(
    device float*         real_buf   [[buffer(0)]],
    device float*         imag_buf   [[buffer(1)]],
    constant FftParams&   params     [[buffer(2)]],
    uint gid [[thread_position_in_grid]]
) {
    uint N = params.length;
    uint stage = params.stage;
    uint half_block = 1u << stage;
    uint block_size = half_block << 1;

    // Determine which FFT row/column this thread belongs to
    uint fft_idx = gid / (N / 2);  // which 1D FFT
    uint local_id = gid % (N / 2); // position within the FFT

    // Base offset into the buffer for this FFT line
    uint base;
    uint elem_stride;
    if (params.is_column == 0) {
        // Row pass: elements are contiguous
        base = fft_idx * N;
        elem_stride = 1;
    } else {
        // Column pass: elements are strided by width
        base = fft_idx; // column index
        elem_stride = params.width;
    }

    // Compute butterfly indices
    uint block_idx = local_id / half_block;
    uint inner_idx = local_id % half_block;
    uint i = block_idx * block_size + inner_idx;
    uint j = i + half_block;

    uint idx_i = base + i * elem_stride;
    uint idx_j = base + j * elem_stride;

    if (idx_i >= params.total_count || idx_j >= params.total_count) return;

    // Twiddle factor: W_N^k = exp(-2*pi*i*k/N) for forward
    float angle = -2.0f * M_PI_F * (float)(block_idx * half_block + inner_idx) / (float)block_size;
    // For the butterfly at this stage, the twiddle exponent is: k / block_size
    // where k = inner_idx
    angle = -2.0f * M_PI_F * (float)inner_idx / (float)block_size;
    if (params.inverse != 0) {
        angle = -angle;
    }

    float tw_re = cos(angle);
    float tw_im = sin(angle);

    float a_re = real_buf[idx_i];
    float a_im = imag_buf[idx_i];
    float b_re = real_buf[idx_j];
    float b_im = imag_buf[idx_j];

    // W * b
    float wb_re = tw_re * b_re - tw_im * b_im;
    float wb_im = tw_re * b_im + tw_im * b_re;

    // Butterfly
    real_buf[idx_i] = a_re + wb_re;
    imag_buf[idx_i] = a_im + wb_im;
    real_buf[idx_j] = a_re - wb_re;
    imag_buf[idx_j] = a_im - wb_im;
}

/// Bit-reversal permutation kernel.
/// Must be run before the butterfly stages.
kernel void fft_bit_reverse(
    device float*         real_buf   [[buffer(0)]],
    device float*         imag_buf   [[buffer(1)]],
    constant FftParams&   params     [[buffer(2)]],
    uint gid [[thread_position_in_grid]]
) {
    uint N = params.length;
    uint log2N = params.log2_length;

    uint fft_idx = gid / N;
    uint local_id = gid % N;

    uint base;
    uint elem_stride;
    if (params.is_column == 0) {
        base = fft_idx * N;
        elem_stride = 1;
    } else {
        base = fft_idx;
        elem_stride = params.width;
    }

    // Compute bit-reversed index
    uint rev = 0;
    uint val = local_id;
    for (uint b = 0; b < log2N; b++) {
        rev = (rev << 1) | (val & 1);
        val >>= 1;
    }

    // Only swap if rev > local_id to avoid double-swapping
    if (rev > local_id) {
        uint idx_a = base + local_id * elem_stride;
        uint idx_b = base + rev * elem_stride;

        if (idx_a < params.total_count && idx_b < params.total_count) {
            float tmp_re = real_buf[idx_a];
            float tmp_im = imag_buf[idx_a];
            real_buf[idx_a] = real_buf[idx_b];
            imag_buf[idx_a] = imag_buf[idx_b];
            real_buf[idx_b] = tmp_re;
            imag_buf[idx_b] = tmp_im;
        }
    }
}

/// Normalization kernel for inverse FFT (divide by N).
kernel void fft_normalize(
    device float*         real_buf   [[buffer(0)]],
    device float*         imag_buf   [[buffer(1)]],
    constant FftParams&   params     [[buffer(2)]],
    uint gid [[thread_position_in_grid]]
) {
    if (gid >= params.total_count) return;
    float inv_n = 1.0f / (float)params.length;
    real_buf[gid] = real_buf[gid] * inv_n;
    imag_buf[gid] = imag_buf[gid] * inv_n;
}