/********************************************************************
 *
 * Research Computing Center
 * The University of Chicago
 * Edelstone Center
 * 6030 S. Ellis Avenue, Room 126
 * Chicago, IL 60637
 * 773-795-2667
 * rcc@uchicago.edu
 *
 * Created: Tue Jul 10 
 ********************************************************************/

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

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

int main( int argc, char *argv[] ) {
	int fft_size;
	fftw_plan p;
	fftw_complex *data;
	ptrdiff_t alloc_local, local_size, local_offset;

	MPI_Init( &argc, &argv );
	fftw_mpi_init();

	if ( argc != 2 ) {
		fprintf(stderr,"Usage: ./a.out N (size of fft NxN)\n");
		exit(1);
	}

	/* allocate local space and initialize arrays */
	fft_size = atoi(argv[1]);
	alloc_local = fftw_mpi_local_size_2d( fft_size, fft_size, 
			MPI_COMM_WORLD, &local_size, &local_offset ); 
	data = fftw_alloc_complex(alloc_local);
	memset( data, 0, alloc_local*sizeof(fftw_complex));
	
	/* create fftw plan, compute fft, and destroy plan */
	p = fftw_mpi_plan_dft_2d( fft_size, fft_size, data, data, 
			MPI_COMM_WORLD, FFTW_FORWARD, FFTW_ESTIMATE );
	fftw_execute(p);
	fftw_destroy_plan(p);

	fftw_free(data);

	MPI_Finalize();
	return 0;
}