#include <stdio.h>
#include <sys/time.h>
#ifndef SINGLE
#include <dfftw_mpi.h>
#else 
#include <sfftw_mpi.h>
#endif

#define space_loop for(x=0;x<lnx;++x)for(y=0;y<ny;++y)for(z=0;z<nz;++z) 
#define recip_loop for(y=0;y<lnyt;++y)for(x=0;x<nx;++x)for(z=0;z<nz;++z)

/*
 * 	nx*ny*nz FFT computed in parallel (MPI FFTW)
 *	
 *	slab decomposed (along x) among numtasks MPI tasks
 *		
 *	lnx, lny, lny		 local nx,ny,nz
 *	lnyt			 local ny after transpose
 *	lxs			 local x start
 *	lyst			 local y start after transpose
 *	lsize			 total local size
 *		
 *	David Skinner	Aug 2001 (dskinner@nersc.gov)
 */
int main(int argc, char **argv)
{
     int 		nn = 100, nx,ny,nz,nit;		/* DIMS */
     int 		numtasks,rank; 			/* MPI */
     int 		lnx, lxs, lnyt, lyst, lsize;	/* LOCAL DIMS */
     int 		it, x, y, z;			/* LOOP VARS */
     fftwnd_mpi_plan 	plan, iplan;			/* PLANS */
     fftw_complex 	*data, *work;			/* VECTORS */
     double		ti,tf,t1,t2,t3;			/* TIMERS */
     double		myerr,err;			/* RESIDUALS */
     struct             timeval tv;                     /* EPOCH */
     char 		cbuf[100];			/* CHAR BUF */

     MPI_Init(&argc,&argv);
     ti = MPI_Wtime();
     if(argc != 3) {
      if(!rank){printf("usage: %s dim_size ncycles\n",argv[0]);}
      if(!rank){printf("\tdim_size = number of points in each dimension\n");}
      if(!rank){printf("\tncycles = number of forward/backward transforms\n");}
      MPI_Finalize();
      fflush(stdout);
      exit(1);
     } else {
      nn = atoi(argv[1]); nit = atoi(argv[2]);
     }
     MPI_Comm_size(MPI_COMM_WORLD, &numtasks);
     MPI_Comm_rank(MPI_COMM_WORLD, &rank);

     gettimeofday(&tv,NULL);
     if(!rank) { printf("MPI_Init done at %d\n", tv.tv_sec); ti=tv.tv_sec; }
     gethostname(cbuf, 99);
     printf("HOSTNAME::%s MPI::%d\n",cbuf,rank);

     nx = ny = nz = nn;

     plan = fftw3d_mpi_create_plan(MPI_COMM_WORLD, nx, ny, nz,
	 FFTW_FORWARD, FFTW_ESTIMATE);

     iplan = fftw3d_mpi_create_plan(MPI_COMM_WORLD, nx, ny, nz,
	 FFTW_BACKWARD, FFTW_ESTIMATE);

/*   slab decomposition */
     fftwnd_mpi_local_sizes(plan, &lnx, &lxs, &lnyt, &lyst, &lsize);

     printf("data_decomp:: %d xslab %d->%d (%d) %d (%.3e MB of %.3e MB total)\n"
	,rank, lxs, lxs+lnx,lnx, lyst,
	lsize*sizeof(fftw_complex)*1.0e-6,
	nn*nn*nn*sizeof(fftw_complex)*1.0e-6);

/*  malloc data and work vectors */
     t1 = MPI_Wtime();
     data = (fftw_complex*) malloc(sizeof(fftw_complex) * lsize);
     if(!data) {printf("malloc failed for data (task = %d)\n",rank);}

     work = (fftw_complex*) malloc(sizeof(fftw_complex) * lsize);
     if(!data) {printf("malloc failed for work (task = %d)\n",rank);}
     t2 = MPI_Wtime();

     if(!rank) {
      printf("time to malloc(%d total = %d tasks * %d l) ... %.3e\n",
	nn*nn*nn,numtasks,lsize,t2-t1);
     }

     space_loop{
        data[(x*ny + y) * nz + z].re = 1.0;
        data[(x*ny + y) * nz + z].im = 0.0;
/*
        printf("input %d %d %d %.3e %.3e\n", x+lxs, y, z,
	 data[(x*ny + y) * nz + z].re,
	 data[(x*ny + y) * nz + z].im);
*/
     }

     for(it = 1; it<=nit; it++) {

      t1 = MPI_Wtime();
      fftwnd_mpi(plan, 1, data, work, FFTW_TRANSPOSED_ORDER);
      fftwnd_mpi(iplan, 1, data, work, FFTW_TRANSPOSED_ORDER);
      t2 = MPI_Wtime();

      err = myerr = 0.0;
      space_loop{
       myerr += (nn*nn*nn-data[(x*ny + y) * nz + z].re)*
		(nn*nn*nn-data[(x*ny + y) * nz + z].re) +
	        (data[(x*ny + y) * nz + z].im)*
	        (data[(x*ny + y) * nz + z].im);
      }

      if(myerr > 0.0) printf("ERROR:: rank=%d myerr=%.3e on iter %d \n",
	rank,myerr,it); 

      if(!rank && it%(nit/10)==0) printf("progress:: %d %% complete\n",
	 100*it/nit);

      space_loop{
        data[(x*ny + y) * nz + z].re = 1.0;
        data[(x*ny + y) * nz + z].im = 0.0;
      }

     }

     free(data);
     free(work);
     fftwnd_mpi_destroy_plan(plan);
     fftwnd_mpi_destroy_plan(iplan);
     MPI_Finalize();

     if(!rank) {
      gettimeofday(&tv,NULL);
      printf("MPI_Finalize done at %d\n", tv.tv_sec); tf=tv.tv_sec;
      printf("total_time:: %.3e seconds\n", tf-ti);
     }
     
}