!---------------------------------------------------------------------------
! pmlib_fourier_fft_fftw.f
!   version: 2D single-core
!---------------------------------------------------------------------------
! <<HEADER>>
!---------------------------------------------------------------------------
SUBROUTINE pmlib_fft( topo_all,topo_in,mesh,topo_out,mesh_fft,fft_seq, & 
                    & ierr,fft_shift)

IMPLICIT NONE

INCLUDE 'fftw3.f'
!---------------------------------------------------------------------------
! Arguments
!---------------------------------------------------------------------------
  TYPE(class_topology_all)                        :: topo_all
  TYPE(class_topology)                            :: topo_in
  TYPE(class_topology)                            :: topo_out
  REAL(MK),DIMENSION(:,:,:),POINTER               :: mesh
  INTEGER, DIMENSION(pmlib_ndim), INTENT(IN)      :: fft_seq
  LOGICAL, OPTIONAL, INTENT(IN)                   :: fft_shift
  COMPLEX(MKC),DIMENSION(:,:,:),POINTER           :: mesh_fft
  INTEGER, INTENT(OUT)                            :: ierr

!---------------------------------------------------------------------------
! Local variables
!---------------------------------------------------------------------------
  CHARACTER(LEN=10)                 :: caller = 'pmlib_fft'
  INTEGER                           :: i,j,k,l,m

  INTEGER, DIMENSION(pmlib_ndim)    :: ncell
  INTEGER, DIMENSION(2*pmlib_ndim)  :: nghost
  REAL(MK), DIMENSION(pmlib_ndim)   :: xmin, xmax, dx 

  INTEGER                           :: nvar
  INTEGER                           :: nfft

  COMPLEX(MKC),DIMENSION(:),POINTER :: fft_in  => NULL()
  COMPLEX(MKC),DIMENSION(:),POINTER :: fft_out => NULL()

  INTEGER                           :: fft_plan

  TYPE(class_topology)              :: topo_send
  TYPE(class_topology)              :: topo_recieve

!---------------------------------------------------------------------------
! Initiate subroutine
!---------------------------------------------------------------------------
  ierr = 0
  nvar = SIZE(mesh,1)

!---------------------------------------------------------------------------
! Move the real mesh to the complex fft mesh
!---------------------------------------------------------------------------
  ncell  = topo_in%ncell
  nghost = topo_in%nghost

  IF( ASSOCIATED(mesh_fft) ) THEN
    DEALLOCATE(mesh_fft,stat=ierr)
    IF (ierr .NE. 0) THEN
      CALL pmlib_write(caller,'Failed to deallocate array.')
      GOTO 9999
    ENDIF
  END IF

  ALLOCATE(mesh_fft(nvar, &
          & 1-nghost(1):ncell(1)+nghost(2), &
          & 1-nghost(3):ncell(2)+nghost(4)), &
          & stat=ierr)
  IF (ierr .NE. 0) THEN
    CALL pmlib_write(caller,'Failed to allocate array.')
    GOTO 9999
  ENDIF

    DO j = 1 - nghost(3), ncell(2) + nghost(4)
      DO i = 1 - nghost(1), ncell(1) + nghost(2)
        DO l = 1,nvar
          mesh_fft(l,i,j) = CMPLX(mesh(l,i,j),0.0_MK,MKC)
        END DO
      END DO
    END DO

!---------------------------------------------------------------------------
! Fourier transform
!---------------------------------------------------------------------------
  DO m = 1,pmlib_ndim
    IF(fft_seq(m) .EQ. 0)THEN
      GOTO 100
    END IF
!---------------------------------------------------------------------------
! Map topology to the fft direction
!---------------------------------------------------------------------------
    IF(m .EQ. 1)THEN
      topo_send = topo_in
    END IF

    IF(fft_seq(m) .EQ. 1)THEN
      topo_recieve = topo_all%xpencil
    ELSEIF(fft_seq(m) .EQ. 2)THEN
      topo_recieve = topo_all%ypencil
    ELSE
        ierr = -1
        CALL pmlib_write(caller, 'fft direction unknown.')
        GOTO 9999
    END IF

    CALL pmlib_mesh_map(topo_send,topo_recieve,2*nvar,ierr)
    CALL pmlib_comm_pack(mesh_fft,ierr)
    CALL pmlib_comm_unpack(topo_recieve,mesh_fft,ierr)
    CALL pmlib_comm_finalise(ierr)

!---------------------------------------------------------------------------
! Do FFTs
!---------------------------------------------------------------------------
    IF(fft_seq(m) .EQ. 1)THEN ! FFT x-pencils

      ncell = topo_recieve%ncell
      nfft  = ncell(1)
      IF(MOD(nfft,2) .NE. 0)THEN
        ierr = -1
        CALL pmlib_write(caller, 'nfft is an odd number.')
        GOTO 9999
      END IF

      ALLOCATE(fft_in(nfft), fft_out(nfft))

! Setup fftw plan
      CALL dfftw_plan_dft_1d( fft_plan, nfft, fft_in, fft_out, & 
                            & FFTW_FORWARD, FFTW_ESTIMATE)

      DO j = 1, ncell(2)
        DO l = 1,nvar

! Store fft pencil
          IF(PRESENT(fft_shift) .AND. fft_shift)THEN
            DO i = 1, nfft/2
              fft_in(i) = mesh_fft(l,i + nfft/2,j)
            END DO
            DO i = nfft/2 + 1, nfft
              fft_in(i) = mesh_fft(l,i - nfft/2,j)
            END DO
          ELSE
            DO i = 1, nfft
              fft_in(i) = mesh_fft(l,i,j)
            END DO
          END IF
! Fourier transform
          CALL dfftw_execute_dft( fft_plan, fft_in, fft_out )

! Store mesh
          DO i = 1, ncell(1)
            mesh_fft(l,i,j) = fft_out(i)
          END DO

       END DO !l
      END DO !j

! Destroy fftw plan
      CALL dfftw_destroy_plan( fft_plan )
      CALL dfftw_cleanup()

    ELSEIF(fft_seq(m) .EQ. 2)THEN ! FFT y-pencils

      ncell = topo_recieve%ncell
      nfft  = ncell(2)
      IF(MOD(nfft,2) .NE. 0)THEN
        ierr = -1
        CALL pmlib_write(caller, 'nfft is an odd number.')
        GOTO 9999
      END IF

      ALLOCATE(fft_in(nfft), fft_out(nfft))

! Setup fftw plan
      CALL dfftw_plan_dft_1d( fft_plan, nfft, fft_in, fft_out, & 
                            & FFTW_FORWARD, FFTW_ESTIMATE)

      DO i = 1, ncell(1)
        DO l = 1,nvar

! Store fft pencil
          IF(PRESENT(fft_shift) .AND. fft_shift)THEN
            DO j = 1, nfft/2
              fft_in(j) = mesh_fft(l,i,j + nfft/2)
            END DO
            DO j = nfft/2 + 1, nfft
              fft_in(j) = mesh_fft(l,i,j - nfft/2)
            END DO
          ELSE
            DO j = 1, nfft
              fft_in(j) = mesh_fft(l,i,j)
            END DO
          END IF
! Fourier transform
          CALL dfftw_execute_dft( fft_plan, fft_in, fft_out )

! Store mesh
          DO j = 1, ncell(2)
            mesh_fft(l,i,j) = fft_out(j)
          END DO

        END DO !l
      END DO !i

! Destroy fftw plan
      CALL dfftw_destroy_plan( fft_plan )
      CALL dfftw_cleanup()

    END IF ! fft direction

    DEALLOCATE(fft_in, fft_out)

!---------------------------------------------------------------------------
! Shift topology reciever to sender
!---------------------------------------------------------------------------
    topo_send = topo_recieve

  END DO ! fft_seq
 100 CONTINUE

!---------------------------------------------------------------------------
! Map to the output topology
!---------------------------------------------------------------------------
  CALL pmlib_mesh_map(topo_send,topo_out,2*nvar,ierr)
  CALL pmlib_comm_pack(mesh_fft,ierr)
  CALL pmlib_comm_unpack(topo_out,mesh_fft,ierr)
  CALL pmlib_comm_finalise(ierr)

!---------------------------------------------------------------------------
! Deallocate local pointers
!---------------------------------------------------------------------------

!---------------------------------------------------------------------------
! Return
!---------------------------------------------------------------------------
 9999 CONTINUE
RETURN
END SUBROUTINE pmlib_fft
