#include "p3dfft.h"
#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>

int main(int argc,char **argv)
{
   double *A,*B,*p,*C;
   int i, j, k, x, y, z, nx, ny, nz;
   int proc_id, nproc, dims[2];
   int memsize[3];
   int istart[3], isize[3], iend[3];
   int fstart[3], fsize[3], fend[3];
   int conf, iproc, jproc, ng[3];
   long int Nglob, Ntot;
   double pi, twopi, sinyz;
   double cdiff, ccdiff, ans, prec;
   double *sinx, *siny, *sinz, factor;
   unsigned char op_f[3]="fft", op_b[3]="tff";

   /* MPI Initializzation */
   MPI_Init(&argc,&argv);
   MPI_Comm_size(MPI_COMM_WORLD,&nproc);
   MPI_Comm_rank(MPI_COMM_WORLD,&proc_id);

   /* Dimension definition */ 
   nx=ny=nz=128; 

   /* Dims create */
   dims[0]=dims[1]=0;
   MPI_Dims_create(nproc,2,dims);
   if(dims[0] > dims[1]) {
     dims[0] = dims[1];
     dims[1] = nproc/dims[0];
    }
   /* Other initializzations */
   pi = atan(1.0)*4.0;
   twopi = 2.0*pi;
   /* Initialize P3DFFT */
   p3dfft_setup(dims,nx,ny,nz,1,memsize);

   MPI_Barrier(MPI_COMM_WORLD);

   if(proc_id == 0) printf("Calling get_dims 1\n");
   conf = 1;
   /* Get dimensions for input array - real numbers, X-pencil shape.
      Note that we are following the Fortran ordering, i.e. 
      the dimension  with stride-1 is X. */
   p3dfft_get_dims(istart,iend,isize,conf);

   if(proc_id == 0) printf("Calling get_dims 2\n");
   conf = 2;
   /* Get dimensions for output array - complex numbers, Z-pencil shape.
      Stride-1 dimension could be X or Z, depending on how the library 
      was compiled (stride1 option) */
   p3dfft_get_dims(fstart,fend,fsize,conf);

   if(proc_id == 0) printf("nx = %d;  ny = %d;   nz = %d\n", nx, ny, nz);

   /* Allocatring */
   if(proc_id == 0) printf("Allocating A\n");
   A = (double *) malloc(sizeof(double) * isize[0]*isize[1]*isize[2]);
   if(proc_id == 0) printf("Allocating B\n");
   B = (double *) malloc(sizeof(double) * fsize[0]*fsize[1]*fsize[2]*2);
   if(proc_id == 0) printf("Allocating C\n");
   C = (double *) malloc(sizeof(double) * isize[0]*isize[1]*isize[2]);

   if(A == NULL) 
     printf("%d: Error allocating array A (%d)\n",proc_id,isize[0]*isize[1]*isize[2]);
   if(B == NULL) 
     printf("%d: Error allocating array B (%d)\n",proc_id,fsize[0]*fsize[1]*fsize[2]*2);
   if(C == NULL) 
     printf("%d: Error allocating array C (%d)\n",proc_id,isize[0]*isize[1]*isize[2]);

   sinx = malloc(sizeof(double)*nx);
   siny = malloc(sizeof(double)*ny);
   sinz = malloc(sizeof(double)*nz);
   /* Initial Solution */
   if(proc_id == 0) printf("Initial solution Step 1 of 2\n");
   for(z=0;z < isize[2];z++){
     sinz[z] = sin((z+istart[2]-1)*twopi/nz);
   }
   for(y=0;y < isize[1];y++){
     siny[y] = sin((y+istart[1]-1)*twopi/ny);
   }
   for(x=0;x < isize[0];x++){
     sinx[x] = sin((x+istart[0]-1)*twopi/nx);
   }
   if(proc_id == 0) printf("Initial solution Step 2 of 2\n");
   p = A;
   for(z=0;z < isize[2];z++){
     for(y=0;y < isize[1];y++) {
       sinyz = siny[y]*sinz[z];
       for(x=0;x < isize[0];x++)
          *p++ = sinx[x]*sinyz;
     }
   }
   if(proc_id == 0) printf("Create Normalizzation Factor\n");
   Ntot = fsize[0]*fsize[1];
   Ntot *= fsize[2]*2;
   Nglob = nx * ny;
   Nglob *= nz;
   factor = 1.0/Nglob;

   MPI_Barrier(MPI_COMM_WORLD);
   if(proc_id == 0) printf("Real to complex DFT upward\n");
   p3dfft_ftran_r2c(A,B,op_f);

   MPI_Barrier(MPI_COMM_WORLD);
   if(proc_id == 0) printf("Normalizzation\n");
   /* Normalizzation */
   for(i=0;i < Ntot;i++){
     B[i] *= factor;
   }

   MPI_Barrier(MPI_COMM_WORLD);
   if(proc_id == 0) printf("Real to complex DFT backward\n");
   p3dfft_btran_c2r(B,C,op_b);

   MPI_Barrier(MPI_COMM_WORLD);
   if(proc_id == 0) printf("P3DFFT cleanup\n");
   p3dfft_clean();
  
   if(proc_id == 0) printf("Check Results\n");
   /* Check results */
   cdiff = 0.0; p = C;
   for(z=0;z < isize[2];z++){
     for(y=0;y < isize[1];y++)  {
        sinyz =siny[y]*sinz[z];
        for(x=0;x < isize[0];x++) {
           ans = sinx[x]*sinyz;
           if(cdiff < fabs(*p - ans))
            cdiff = fabs(*p - ans);
	    p++;
         }
     }
   }
   MPI_Reduce(&cdiff,&ccdiff,1,MPI_DOUBLE,MPI_MAX,0,MPI_COMM_WORLD);

  if(proc_id == 0) {
    prec = 1.0e-7;

    if(ccdiff > prec * Nglob*0.25)
      printf("Results are incorrect\n");
    else
      printf("Results are correct\n");

    printf("max diff =%g\n",ccdiff);
  }

   if(proc_id == 0) printf("MPI Finalize\n");
  MPI_Finalize();

}