#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "mpi.h"
#define am(p,q,r) *(a+2*nx*nx*p+2*nx*q+r)
int main(int argc, char **argv )
{
/* 
   3-D binary radix FFT using message passing. Data are
   in slabs. Given "size" processors, there are "size"
   slabs containing NZ/size planes each. Complex
   element A(i,j,k) is in slab i mod (NZ/size) and has 
   location &A(i%(NZ/size),j,k) = &A(mod(i,NZ/size),j,k). 

   From Sections 5.8 and 5.9 of Arbenz and Petersen,
   "Intro. to Parallel Computing," Oxford Univ. Press, 2004.

                                W. Petersen, 22 Oct. 2004
*/

   int MAX_X=128, MAX_Z=128;
   int ierr,master,rank,size,nz,nx,i,j,ijk,ip,k,offset; 
   static float seed;
   float *a,*b,*w,*buff,*acopy,sign,err,fnn;
   float ggl();
   void cffti(int nx, float *w);
   void Checkres(float *a, float *acopy, int nz, int nx);
   void FFT3D(float *a,float *w,float sign,int nz,int nx);
   MPI_Status stat;

   MPI_Init(&argc,&argv);
   MPI_Comm_size(MPI_COMM_WORLD, &size);
   MPI_Comm_rank(MPI_COMM_WORLD, &rank);
   master = 0;

   nx     = MAX_X;
   nz     = nx/size;
   a      = (float *) malloc(2*nz*nx*nx*sizeof(float));
   acopy  = (float *) malloc(2*nx*nx*nx*sizeof(float));
   w      = (float *) malloc(2*nx*sizeof(float));
   buff   = (float *) malloc(2*nx*nx*nz*sizeof(float));
/* initialize sine/cosine tables */
   cffti(nx,w);
   if(rank==master){
/* now send section of A to each proccessor */
     seed = 331.0;
     for(ip=0;ip<size;ip++){
        if(ip==0){
           ijk = 0;
           for(k=0;k<nz;k++){
              for(j=0;j<nx;j++){
                 for(i=0;i<nx;i++){
                    am(k,j,2*i)    = ggl(&seed); 
                    am(k,j,2*i+1)  = ggl(&seed); 
                    acopy[ijk]   = am(k,j,2*i);
                    acopy[ijk+1] = am(k,j,2*i+1);
                    ijk += 2;
                 }
              }
           }
        } else {
           ijk = 0; offset=2*ip*nx*nx*nz;
           for(k=0;k<nz;k++){
              for(j=0;j<nx;j++){
                 for(i=0;i<nx;i++){ 
                     buff[ijk]   = ggl(&seed); 
                     buff[ijk+1] = ggl(&seed); 
                     acopy[offset+ijk]   = buff[ijk];
                     acopy[offset+ijk+1] = buff[ijk+1];
                     ijk += 2;
                 }
              }
           }
        }
        if(ip>0){
           ierr=MPI_Send(buff,2*nx*nx*nz,MPI_FLOAT,ip,ip,
                         MPI_COMM_WORLD);
        }
     }
   } else { /* slave parts */
      ierr=MPI_Recv(buff,2*nx*nx*nz,MPI_FLOAT,master,
                    MPI_ANY_TAG,MPI_COMM_WORLD,&stat);
      if(stat.MPI_TAG!=0){
         ijk = 0;
         for(k=0;k<nz;k++){
            for(j=0;j<nx;j++){
               for(i=0;i<2*nx;i++){
                  am(k,j,i) = buff[ijk++];
               }
            }
         }
      }
   }
   sign = 1.0;
   FFT3D(a,w,sign,nz,nx);
   sign = -1.0;
   FFT3D(a,w,sign,nz,nx); 

   Checkres(a,acopy,nz,nx); 
   free(a); free(acopy); free(w); free(buff); 
   MPI_Finalize();
}
void FFT3D(float *a,float *w,float sign,int nz,int nx)
{
   int i,j,k,offset,rank;
   static int nfirst=-1;
   static float *pw;
   float *pa;
   void Xpose();
   void cfft2();

   MPI_Comm_rank(MPI_COMM_WORLD,&rank);
   if(nfirst!=nx){
      if(nfirst>0) free(pw);
      pw     = (float *) malloc(2*nx*sizeof(float));
      nfirst = nx;
   }
/* X-direction */
   for(k=0;k<nz;k++){
      offset = 2*k*nx*nx;
      for(j=0;j<nx;j++){
         pa = a + offset + 2*nx*j;
         cfft2(nx,pa,w,sign);
      }
   } 
/* Y-direction  */
   for(k=0;k<nz;k++){
      for(i=0;i<nx;i++){
         offset = 2*k*nx*nx+2*i;
         for(j=0;j<nx;j++){
            *(pw+2*j)   = *(a+2*j*nx+offset);
            *(pw+2*j+1) = *(a+2*j*nx+1+offset);
         }
         cfft2(nx,pw,w,sign); 
         for(j=0;j<nx;j++){
            *(a+2*j*nx+offset)   = *(pw+2*j);
            *(a+2*j*nx+1+offset) = *(pw+2*j+1);
         }
      }
   }
/* Z-direction */
   Xpose(a,nz,nx);
   for(k=0;k<nz;k++){
      offset = 2*k*nx*nx;
      for(j=0;j<nx;j++){
         pa = a + offset + 2*nx*j;
         cfft2(nx,pa,w,sign);
      }
   }
   Xpose(a,nz,nx);
}

void Checkres(float *a,float *acopy,int nz,int nx)
{
   int i,ierr,ijk,is,j,k,offset,rank,size,sor;
   float err,fnm3,*buff;
   MPI_Status stat;

   MPI_Comm_rank(MPI_COMM_WORLD, &rank);
   MPI_Comm_size(MPI_COMM_WORLD, &size);

   buff = (float *) malloc(2*nz*nx*nx*sizeof(float));

   fnm3 = 1.0/((float) (nx*nx*nx)); 
   if(rank>0){
      ierr=MPI_Send(a,2*nz*nx*nx,MPI_FLOAT,0,0,
                    MPI_COMM_WORLD);
   } else { /* master */
/* rank > 0 part of check */
      for(is=1;is<size;is++){
         ierr = MPI_Recv(buff,2*nz*nx*nx,MPI_FLOAT,
                MPI_ANY_SOURCE,MPI_ANY_TAG,MPI_COMM_WORLD,
                &stat);
         sor  = stat.MPI_SOURCE; ijk = 0;
         err  = 0.0; offset = 2*nz*nx*nx*sor;
         for(k=0;k<nz;k++){
            for(j=0;j<nx;j++){
               for(i=0;i<2*nx;i++){
                  err += (acopy[offset+ijk]-fnm3*buff[ijk])*
                         (acopy[offset+ijk]-fnm3*buff[ijk]);
                  ijk++;
               }
            }
         }
         err = sqrt(fnm3*err);
         printf(" error from cpu %d: N=%d, err=%e\n",sor,nx,err);
      }
/* rank = 0 part of check */
      err  = 0.0; ijk = 0;
      for(k=0;k<nz;k++){
         for(j=0;j<nx;j++){
            for(i=0;i<2*nx;i++){
               err += (acopy[ijk]-fnm3*a[ijk])*
                      (acopy[ijk]-fnm3*a[ijk]);
               ijk++;
            }
         }
      }
      err = sqrt(fnm3*err);
      printf(" error from cpu %d: N=%d, err=%e\n",rank,nx,err);
      ierr = fflush(stdout);
   }
   free(buff);
}
void Xpose(float *a, int nz, int nx) {
  float t0,t1;
  static float *buf_io; 
  int i, ijk, j, js, k, step, n2, np, offset;
  static int init=-1;
  int size, rank, other;
  MPI_Status stat;
 
  MPI_Comm_size(MPI_COMM_WORLD, &size);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
 
  /* number of local planes of 3D array */
  n2 = 2*nz; 
  np = 2*nx*nx; 

  if(init!=nx){
     if(init>0) free(buf_io);
     buf_io  = (float *) malloc(nz*n2*nx*sizeof(float));
     init    = nx;
  }
  /* local transpose of first block (in-place) */
  for(j = 0; j < nx; j++){
     offset = j*2*nx + rank*n2;
     for(k = 0; k < nz; k ++){
        for(i = 0; i < k; i++) {
           t0 = a[offset + i*np + k*2];
           t1 = a[offset + i*np + k*2+1];
           a[offset+i*np+k*2]   = a[offset+k*np+2*i];
           a[offset+i*np+k*2+1] = a[offset+k*np+2*i+1];
           a[offset+k*np+2*i]   = t0;
           a[offset+k*np+2*i+1] = t1;
        }
     }
  }

  /* size-1 communication steps */
  for (step = 1; step < size; step ++) {
    other = rank ^ step;

    /* fill send buffer */
    ijk = 0;
    for(j=0;j<nx;j++){
       for(k=0;k<nz;k++){
          offset = j*2*nx + other*n2 + k*np;  
          for(i=0;i<n2;i++){
             buf_io[ijk++] = a[offset + i];
          }
       } 
    } 
    /* exchange data */
    MPI_Sendrecv_replace(buf_io,n2*nz*nx,MPI_FLOAT,
                   other,rank,other,other,
                   MPI_COMM_WORLD,&stat);

    /* write back recv buffer in transposed order */
    ijk = 0;
    for(j=0;j<nx;j++){
       offset = j*2*nx + other*n2; 
       for(k=0;k<nz;k++){
          for(i=0;i<nz;i++){
             a[offset+i*np+2*k]   = buf_io[ijk];
             a[offset+i*np+2*k+1] = buf_io[ijk+1];
             ijk += 2;
          }
       }
    }
  }
}
void cffti(n,w)
int n;
float w[][2];
{
   int i;
   float pi,arg,t;

   pi = 3.141592653589793;
   t  = 2.0*pi/((float) n);
   for(i=0;i<n/2;i++){
      arg     = t*((float) i);
      w[i][0] = cos(arg);
      w[i][1] = sin(arg);
   }
}
void cfft2(n,x,w,sign)
int n;
float sign, x[][2], w[][2];
{
   int n2,m,j,mj,p2,p3,p4,BK;
   void step0(), step1(), step2();

   m  = (int) (log((float) n)/log(1.99));
/* plain vanilla: */
   BK = 0; 

   mj = 1;
   n2 = n/2;

   for(j=0;j<m;j++){
      if(j < (m+1)/2){
         p2 = n2/mj;
         step0(n,mj,&x[0][0],&x[p2][0],w,sign);
      }
      else{
         p2 = n2/mj;
         p3 = mj;
         p4 = p2+mj;
         if(j < (m/2+BK)){
            step1(n,mj,&x[0][0],&x[p2][0],
                  &x[p3][0],&x[p4][0],w,sign);
         } 
         else{
            step2(n,mj,&x[0][0],&x[p2][0],
                  &x[p3][0],&x[p4][0],w,sign);
         }
      }
      mj = 2*mj;
   }
}
void step0(n,mj,a,b,w,sign)
int n,mj;
float a[][2],b[][2],w[][2];
float sign;
{
   float wkr,wku,wambr,wambu;
   int i,k,kw,lj,ii,ij;

   lj = n/(2*mj);
   ij = n/mj;
   for(i=0;i<mj;i++){
      ii = i*ij;
      if(sign > 0.){
#pragma ivdep
         for(k=0;k<lj;k++){
             kw    = k*mj;
             wkr   = w[kw][0];
             wku   = w[kw][1];
             wambr = wkr*(a[ii+k][0]-b[ii+k][0]) 
                   - wku*(a[ii+k][1]-b[ii+k][1]);
             wambu = wku*(a[ii+k][0]-b[ii+k][0]) 
                   + wkr*(a[ii+k][1]-b[ii+k][1]);
             a[ii+k][0] = a[ii+k][0]+b[ii+k][0];
             a[ii+k][1] = a[ii+k][1]+b[ii+k][1];
             b[ii+k][0] = wambr;
             b[ii+k][1] = wambu;
         } 
      } else {
#pragma ivdep
         for(k=0;k<lj;k++){
             kw    = k*mj;
             wkr   = w[kw][0];
             wku   = -w[kw][1];
             wambr = wkr*(a[ii+k][0]-b[ii+k][0]) 
                   - wku*(a[ii+k][1]-b[ii+k][1]);
             wambu = wku*(a[ii+k][0]-b[ii+k][0]) 
                   + wkr*(a[ii+k][1]-b[ii+k][1]);
             a[ii+k][0] = a[ii+k][0]+b[ii+k][0];
             a[ii+k][1] = a[ii+k][1]+b[ii+k][1];
             b[ii+k][0] = wambr;
             b[ii+k][1] = wambu;
         } 
      }
   }
}
void step1(n,mj,a,b,c,d,w,sign)
int n,mj;
float a[][2],b[][2],c[][2],d[][2],w[][2];
float sign;
{
   float wkr,wku,wambr,wambu,wcmdr,wcmdu;
   int mj2,i,j,k,kw,lj;

   mj2 = 2*mj;
   lj  = n/mj2;

   for(j=0;j<mj;j+=n/mj){
      if(sign > 0.){
         for(i=j;i<n;i+=mj2){
#pragma ivdep
            for(k=0;k<lj;k++){
                kw    = k*mj;
                wkr   = w[kw][0];
                wku   = w[kw][1];
                wambr = wkr*(a[i+k][0]-b[i+k][0]) 
                      - wku*(a[i+k][1]-b[i+k][1]);
                wambu = wku*(a[i+k][0]-b[i+k][0]) 
                      + wkr*(a[i+k][1]-b[i+k][1]);
                a[i+k][0] = a[i+k][0]+b[i+k][0];
                a[i+k][1] = a[i+k][1]+b[i+k][1];
                b[i+k][0] = c[i+k][0]+d[i+k][0];
                b[i+k][1] = c[i+k][1]+d[i+k][1];
                wcmdr = wkr*(c[i+k][0]-d[i+k][0]) 
                      - wku*(c[i+k][1]-d[i+k][1]);
                wcmdu = wku*(c[i+k][0]-d[i+k][0]) 
                      + wkr*(c[i+k][1]-d[i+k][1]);
                c[i+k][0] = wambr;
                c[i+k][1] = wambu;
                d[i+k][0] = wcmdr;
                d[i+k][1] = wcmdu;
            } 
         } 
      } else {
         for(i=j;i<n;i+=mj2){
#pragma ivdep
            for(k=0;k<lj;k++){
                kw    = k*mj;
                wkr   = w[kw][0];
                wku   = -w[kw][1];
                wambr = wkr*(a[i+k][0]-b[i+k][0]) 
                      - wku*(a[i+k][1]-b[i+k][1]);
                wambu = wku*(a[i+k][0]-b[i+k][0]) 
                      + wkr*(a[i+k][1]-b[i+k][1]);
                a[i+k][0] = a[i+k][0]+b[i+k][0];
                a[i+k][1] = a[i+k][1]+b[i+k][1];
                b[i+k][0] = c[i+k][0]+d[i+k][0];
                b[i+k][1] = c[i+k][1]+d[i+k][1];
                wcmdr = wkr*(c[i+k][0]-d[i+k][0]) 
                      - wku*(c[i+k][1]-d[i+k][1]);
                wcmdu = wku*(c[i+k][0]-d[i+k][0]) 
                      + wkr*(c[i+k][1]-d[i+k][1]);
                c[i+k][0] = wambr;
                c[i+k][1] = wambu;
                d[i+k][0] = wcmdr;
                d[i+k][1] = wcmdu;
            }
         }
      }
   }
}
void step2(n,mj,a,b,c,d,w,sign)
int n,mj;
float a[][2],b[][2],c[][2],d[][2],w[][2];
float sign;
{
   float wkr,wku,wambr,wambu,wcmdr,wcmdu;
   int mj2,i,j,k,kw,lj,ii;

   mj2 = 2*mj;
   lj  = n/mj2;

   for(k=0;k<lj;k++){
      kw    = k*mj;
      wkr   = w[kw][0];
      if(sign > 0.){
         wku   = w[kw][1];
      } else {
         wku   = -w[kw][1];
      }
      for(i=0;i<lj;i++){
         ii = i*mj2;
#pragma ivdep
         for(j=k;j<mj;j+=n/mj){
             wambr = wkr*(a[ii+j][0]-b[ii+j][0]) 
                   - wku*(a[ii+j][1]-b[ii+j][1]);
             wambu = wku*(a[ii+j][0]-b[ii+j][0]) 
                   + wkr*(a[ii+j][1]-b[ii+j][1]);
             a[ii+j][0] = a[ii+j][0]+b[ii+j][0];
             a[ii+j][1] = a[ii+j][1]+b[ii+j][1];
             b[ii+j][0] = c[ii+j][0]+d[ii+j][0];
             b[ii+j][1] = c[ii+j][1]+d[ii+j][1];
             wcmdr = wkr*(c[ii+j][0]-d[ii+j][0]) 
                   - wku*(c[ii+j][1]-d[ii+j][1]);
             wcmdu = wku*(c[ii+j][0]-d[ii+j][0]) 
                   + wkr*(c[ii+j][1]-d[ii+j][1]);
             c[ii+j][0] = wambr;
             c[ii+j][1] = wambu;
             d[ii+j][0] = wcmdr;
             d[ii+j][1] = wcmdu;
         } 
      } 
   } 
}
#include <math.h>
float ggl(float *ds)
{

/* generate u(0,1) distributed random numbers. 
   Seed ds must be saved between calls. ggl is 
   essentially the same as the IMSL routine RNUM. 

   W. Petersen and M. Troyer, 24 Oct. 2002, ETHZ: 
   a modification of a fortran version from 
   I. Vattulainen, Tampere Univ. of Technology, 
   Finland, 1992 */

   double t,d2=0.2147483647e10;
   t   = (float) *ds;
   t   = fmod(0.16807e5*t,d2);
   *ds = (float) t;
   return((float) ((t-1.0e0)/(d2-1.0e0)));
}