#include <fftw3-mpi.h>
#include <math.h>
#include <stdio.h>

int main(int argc, char **argv)
{
    const ptrdiff_t rows = 8, columns = 4;
    ptrdiff_t local_rows, first_row;
    int rank, size, valid = 1, all_valid;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    fftw_mpi_init();
    ptrdiff_t allocated = fftw_mpi_local_size_2d(rows, columns, MPI_COMM_WORLD,
                                               &local_rows, &first_row);
    fftw_complex *data = fftw_alloc_complex(allocated > 0 ? allocated : 1);
    if (!data) MPI_Abort(MPI_COMM_WORLD, 1);
    fftw_plan plan = fftw_mpi_plan_dft_2d(rows, columns, data, data,
                                         MPI_COMM_WORLD, FFTW_FORWARD, FFTW_ESTIMATE);
    if (!plan) MPI_Abort(MPI_COMM_WORLD, 1);
    for (ptrdiff_t i = 0; i < local_rows; ++i)
        for (ptrdiff_t j = 0; j < columns; ++j) {
            ptrdiff_t index = i * columns + j;
            data[index][0] = (first_row + i == 0 && j == 0) ? 1.0 : 0.0;
            data[index][1] = 0.0;
        }
    fftw_execute(plan);
    for (ptrdiff_t i = 0; i < local_rows * columns; ++i)
        if (fabs(data[i][0] - 1.0) > 1e-12 || fabs(data[i][1]) > 1e-12) valid = 0;
    MPI_Allreduce(&valid, &all_valid, 1, MPI_INT, MPI_MIN, MPI_COMM_WORLD);
    if (rank == 0 && all_valid) printf("PASS: MPI transform with %d processes\n", size);
    fftw_destroy_plan(plan);
    fftw_free(data);
    fftw_mpi_cleanup();
    MPI_Finalize();
    return all_valid ? 0 : 1;
}