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

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

#include <fftw.h>
#include <fftw_mpi.h>
#include <mpi.h>

int main( int argc, char *argv[] ) {
	int fft_size;
	fftw_mpi_plan p;
	fftw_complex *in, *out;
	int local_x, local_x_start;
	int local_ny_after_transpose;
	int local_y_start_after_transpose;
	int total_local_size;

	MPI_Init( &argc, &argv );

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

	fft_size = atoi(argv[1]);

	/* create fftw plan and compute local storage requirements */
	p = fftw_mpi_create_plan( MPI_COMM_WORLD, fft_size,
			FFTW_FORWARD, FFTW_ESTIMATE );

	fftw_mpi_local_sizes( p, &local_x, &local_x_start,
		&local_ny_after_transpose, &local_y_start_after_transpose,
		&total_local_size );

	/* allocate space and initialize real array to transform */
	in = (fftw_complex *)malloc( sizeof(fftw_complex)*total_local_size );
	out = (fftw_complex*)malloc( sizeof(fftw_complex)*total_local_size );

	if ( in == NULL || out == NULL ) {
		fprintf(stderr,"Unable to allocate space for complex arrays!\n");
		exit(1);
	}
	memset( in, 0, total_local_size*sizeof(fftw_complex) );
	
	/* create fftw plan, compute fft, and destroy plan */
	fftw_mpi( p, 1, in, out );
	fftw_mpi_destroy_plan(p);

	free(in);
	free(out);

	MPI_Finalize();
	return 0;
}