use rustfft::num_complex::Complex;
use rustfft::FftPlanner;

/// 3D FFT on a cubic grid of size N×N×N.
///
/// Data is stored in z-major order: index = z*N*N + y*N + x.
/// The FFT is decomposed into successive 1D FFTs along each axis.
pub fn fft3d(data: &mut [Complex<f64>], n: usize) {
    assert_eq!(data.len(), n * n * n);
    let mut planner = FftPlanner::new();
    let fft = planner.plan_fft_forward(n);
    let mut scratch = vec![Complex::new(0.0, 0.0); fft.get_inplace_scratch_len()];

    // FFT along X (innermost axis)
    for z in 0..n {
        for y in 0..n {
            let offset = z * n * n + y * n;
            fft.process_with_scratch(&mut data[offset..offset + n], &mut scratch);
        }
    }

    // FFT along Y — need to gather/scatter strided data
    let mut row = vec![Complex::new(0.0, 0.0); n];
    for z in 0..n {
        for x in 0..n {
            // Gather Y-column
            for y in 0..n {
                row[y] = data[z * n * n + y * n + x];
            }
            fft.process_with_scratch(&mut row, &mut scratch);
            // Scatter back
            for y in 0..n {
                data[z * n * n + y * n + x] = row[y];
            }
        }
    }

    // FFT along Z
    for y in 0..n {
        for x in 0..n {
            // Gather Z-column
            for z in 0..n {
                row[z] = data[z * n * n + y * n + x];
            }
            fft.process_with_scratch(&mut row, &mut scratch);
            for z in 0..n {
                data[z * n * n + y * n + x] = row[z];
            }
        }
    }
}

/// 3D inverse FFT.
pub fn ifft3d(data: &mut [Complex<f64>], n: usize) {
    assert_eq!(data.len(), n * n * n);
    let mut planner = FftPlanner::new();
    let fft = planner.plan_fft_inverse(n);
    let mut scratch = vec![Complex::new(0.0, 0.0); fft.get_inplace_scratch_len()];
    let norm = 1.0 / (n * n * n) as f64;

    // IFFT along X
    for z in 0..n {
        for y in 0..n {
            let offset = z * n * n + y * n;
            fft.process_with_scratch(&mut data[offset..offset + n], &mut scratch);
        }
    }

    // IFFT along Y
    let mut row = vec![Complex::new(0.0, 0.0); n];
    for z in 0..n {
        for x in 0..n {
            for y in 0..n {
                row[y] = data[z * n * n + y * n + x];
            }
            fft.process_with_scratch(&mut row, &mut scratch);
            for y in 0..n {
                data[z * n * n + y * n + x] = row[y];
            }
        }
    }

    // IFFT along Z
    for y in 0..n {
        for x in 0..n {
            for z in 0..n {
                row[z] = data[z * n * n + y * n + x];
            }
            fft.process_with_scratch(&mut row, &mut scratch);
            for z in 0..n {
                data[z * n * n + y * n + x] = row[z];
            }
        }
    }

    // Normalize
    for val in data.iter_mut() {
        *val *= norm;
    }
}
