! cubic_mpi_3dfft testing routine
! wrapper routines to FFT a cubically decomposed data set using the slab decomposed fftw 2.1.5 mpi fft library
! Hugh Merz :: Mon Aug 21 12:26:30 EDT 2006
! compile with: mpif77 -L$MCKENZIE_FFTW_LIB_PATH -I$MCKENZIE_FFTW_INC_PATH -lsrfftw_mpi -lsrfftw -lsfftw_mpi -lsfftw -lm -ldl 3dfft.f90 
implicit none
include 'mpif.h'

integer, parameter :: nodes_dim = 2  !nodes / dimension
integer, parameter :: nc_node_dim=32 !cells / node / dimension

integer, parameter :: nodes_slab = nodes_dim * nodes_dim 
integer, parameter :: nodes = nodes_slab * nodes_dim
integer, parameter :: nc_dim = nc_node_dim * nodes_dim
integer, parameter :: nc_slab = nc_dim / nodes 

logical :: firstfftw

integer(8) :: plan,iplan
integer(4) :: i,j
integer(4), dimension(0:nodes_dim-1,0:nodes_dim-1) :: slab_neighbor
integer(4), dimension(6) :: cart_neighbor
integer(4), dimension(3) :: slab_coord, cart_coords
integer(4) :: slab_rank, nodes_returned, rank, ierr

real(4), dimension(nc_node_dim,nc_node_dim,nc_node_dim) :: cube !real data
real(4), dimension(nc_dim+2,nc_dim,nc_slab) :: slab,slab_work   !fourier space data returned in slab
real(4), dimension(nc_node_dim,nc_node_dim,nc_slab,0:nodes_slab-1)      :: recv_cube

common /rarray/ cube, slab, slab_work, recv_cube

!! check to see if parameters are set properly

if (mod(nc_node_dim,nodes_dim**2) /= 0) then
  print *,'nc_node_dim has to evenly divide by nodes_dim**2'
  print *,'check parameters and recompile'
  stop
endif

!! set up global mpi communicator

call mpi_init(ierr)
call mpi_comm_size(mpi_comm_world,nodes_returned,ierr)
  if (nodes_returned /= nodes ) then
    write(*,*) 'cubepm compiled for a different number of nodes'
    write(*,*) 'mpirun nodes=',nodes_returned,'cubepm nodes=',nodes
    call mpi_abort(mpi_comm_world,ierr,ierr)
  endif
call mpi_comm_rank(mpi_comm_world,rank,ierr)


!! calculate coordinates within slab for cube processes

slab_coord(3) = rank / nodes_slab
slab_rank = rank - slab_coord(3) * nodes_slab
slab_coord(2) = slab_rank / nodes_dim
slab_coord(1) = slab_rank - slab_coord(2) * nodes_dim
do j = 0, nodes_dim - 1
  do i = 0, nodes_dim - 1
    slab_neighbor(i,j) = i + j * nodes_dim + slab_coord(3) &
                       * nodes_slab
  enddo
enddo

!! initialize arrays

cube=1.
slab=0.
slab_work=0.
recv_cube=0.

firstfftw=.true.

print *,'finished init, sum cube=',sum(cube)

call cubic_mpi_3dfft(1)

print *,'finished forward fft, slab(1,1,1)=',slab(1,1,1)

call cubic_mpi_3dfft(-1)

print *,'finished backward fft, sum cube=',sum(cube)

contains

!*******************************************************
! fftw3ds.f90 - cubic decomposition subroutines for fftw

  subroutine pack_slab
!! pack cubic data into slab decomposition for fftw transform
    implicit none
    include 'mpif.h'

    integer(4) :: i,j,k,i0,j0,i1,j1,k1
    integer(4) :: slab_slice,num_elements,tag,rtag

    integer(4), dimension(2*nodes_dim**2) :: requests
    integer(4), dimension(MPI_STATUS_SIZE,2*nodes_dim**2) :: wait_status

    num_elements = nc_node_dim * nc_node_dim * nc_slab

!! swap data

    do j = 0, nodes_dim - 1
      do i = 0, nodes_dim - 1
        slab_slice = i + j * nodes_dim
        tag = rank**2
        rtag= slab_neighbor(i,j)**2
        call mpi_isend(cube(1,1,slab_slice*nc_slab + 1), num_elements, &
                       mpi_real, slab_neighbor(i,j), tag, mpi_comm_world, &
                       requests(slab_slice+1),ierr)
        call mpi_irecv(recv_cube(1,1,1,slab_slice), &
                       num_elements, mpi_real, slab_neighbor(i,j),rtag, &
                       mpi_comm_world, requests(slab_slice+1+nodes_dim**2), &
                       ierr)
      enddo
    enddo

    call mpi_waitall(2*nodes_dim**2, requests, wait_status, ierr)

!! place data in the slab

    do j = 0, nodes_dim - 1
      j0 = j * nc_node_dim + 1
      j1 = (j + 1) * nc_node_dim
      do i = 0, nodes_dim - 1
        i0 = i * nc_node_dim + 1
        i1 = (i + 1) * nc_node_dim
        slab_slice = i + j * nodes_dim
        slab(i0:i1,j0:j1,:) = recv_cube(:,:,:,slab_slice)
      enddo
    enddo

  end subroutine pack_slab

  subroutine unpack_slab 
!! unpack slab data into cubic decomposition following fftw transform
    implicit none
    include 'mpif.h'
    
    integer(4) :: i,j,k,i0,j0,i1,j1,k1
    integer(4) :: slab_slice,num_elements,tag,rtag
    integer(4), dimension(2*nodes_dim**2) :: requests
    integer(4), dimension(MPI_STATUS_SIZE,2*nodes_dim**2) :: wait_status

!! place data in the recv_cube buffer

    do j = 0, nodes_dim - 1
      j0 = j * nc_node_dim + 1
      j1 = (j + 1) * nc_node_dim
      do i = 0, nodes_dim - 1
        i0 = i * nc_node_dim + 1
        i1 = (i + 1) * nc_node_dim
        slab_slice = i + j * nodes_dim
        recv_cube(:,:,:,slab_slice) = slab(i0:i1,j0:j1,:)
      enddo
    enddo

    num_elements = nc_node_dim * nc_node_dim * nc_slab
 
!! swap data

   do j = 0, nodes_dim - 1
      do i = 0, nodes_dim - 1
        slab_slice = i + j * nodes_dim
        tag  = rank**2
        rtag = slab_neighbor(i,j)**2
        call mpi_isend(recv_cube(1,1,1,slab_slice), num_elements, &
                       mpi_real, slab_neighbor(i,j), tag, mpi_comm_world, &
                       requests(slab_slice+1),ierr)
        call mpi_irecv(cube(1,1,slab_slice * nc_slab +1), &
                       num_elements, mpi_real, slab_neighbor(i,j), rtag, &
                       mpi_comm_world, requests(slab_slice+1+nodes_dim**2), &
                       ierr)
      enddo
    enddo

    call mpi_waitall(2*nodes_dim**2,requests, wait_status, ierr)

  end subroutine unpack_slab

  subroutine cubic_mpi_3dfft(command)
!! calculate fftw transform
!! 0 ends fftw subprogram, 1 starts forward fft, -1 starts backwards
    implicit none
    include 'mpif.h'
    include 'fftw_f77.i'

    integer(4), parameter :: order=FFTW_NORMAL_ORDER ! FFTW_TRANSPOSED_ORDER

    integer(4) :: i
    integer(4) :: command 

! initialize plan variables for fftw

    if (firstfftw) then
      call rfftw3d_f77_mpi_create_plan(plan,mpi_comm_world,nc_dim, &
            nc_dim,nc_dim, FFTW_REAL_TO_COMPLEX, FFTW_MEASURE)
      call rfftw3d_f77_mpi_create_plan(iplan,mpi_comm_world,nc_dim, &
            nc_dim,nc_dim, FFTW_COMPLEX_TO_REAL, FFTW_MEASURE)
      firstfftw=.false.
    endif

! start FFT

    if (command /= 0) then

!! call pack routine if we are going forward
  
      if (command > 0) call pack_slab

      if (command > 0) then
        call rfftwnd_f77_mpi(plan,1,slab,slab_work,1,order)
      else
        call rfftwnd_f77_mpi(iplan,1,slab,slab_work,1,order)
        slab=slab/real(nc_dim*nc_dim*nc_dim)
      endif

!! unpack the slab data if we are going backward
 
      if (command < 0) call unpack_slab
 
    else

! if command = 0 we delete the plans
  
      call rfftwnd_f77_mpi_destroy_plan(iplan)
      call rfftwnd_f77_mpi_destroy_plan(plan)
    endif
  
  end subroutine cubic_mpi_3dfft 

end
