!----------------------------------------------------------------------!
! pmlib_poisson_solve_fftw.f
!   version: 2D single-core
!----------------------------------------------------------------------!
! <<HEADER>>
!----------------------------------------------------------------------!
SUBROUTINE pmlib_poisson_solve(patch,topo_all,mesh,ierr,reg_vort)

USE pmlib_mod_patch
USE pmlib_mod_topology
USE pmlib_mod_mesh
USE pmlib_mod_fourier
USE pmlib_mod_communication

IMPLICIT NONE

INCLUDE 'fftw3.f'
!----------------------------------------------------------------------!
! Arguments
!----------------------------------------------------------------------!
  TYPE(class_patch)                             :: patch
  TYPE(class_topology_all)                      :: topo_all
  TYPE(class_mesh)                              :: mesh
  INTEGER, INTENT(OUT)                          :: ierr
  LOGICAL, OPTIONAL, INTENT(IN)                 :: reg_vort

!----------------------------------------------------------------------!
! Local variables
!----------------------------------------------------------------------!
  CHARACTER(LEN=20)                 :: caller = 'pmlib_poisson_solve'
  REAL(MK)                          :: PI

  INTEGER                           :: i,j,l

  REAL(MK)                          :: px, py
  REAL(MK)                          :: rho

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

  INTEGER                           :: i_ptc, j_ptc
  REAL(MK), DIMENSION(pmlib_ndim)   :: l_ptc, ks, dk, kmin
  REAL(MK)                          :: ki, kj, s
  REAL(MK)                          :: eps
  REAL(MK), DIMENSION(pmlib_ndim)   :: xmin_ptc, xmax_ptc
  INTEGER, DIMENSION(pmlib_ndim)    :: bc_ptc

  REAL(MK)                          :: c1, c2, c3, c4, c5

  INTEGER, DIMENSION(pmlib_ndim)    :: fft_seq, ifft_seq
  LOGICAL                           :: ifft_vort

  COMPLEX(MKC)                      :: G, Kx, Ky
  COMPLEX(MKC)                      :: zeta_fft

  INTEGER                           :: nfft

  COMPLEX(MKC),DIMENSION(:,:,:),POINTER :: vel_fft, vort_fft 

  COMPLEX(MKC),DIMENSION(:),POINTER :: fft1_in   => NULL()
  COMPLEX(MKC),DIMENSION(:),POINTER :: fft1_out  => NULL()
  COMPLEX(MKC),DIMENSION(:),POINTER :: ifft1_in  => NULL()
  COMPLEX(MKC),DIMENSION(:),POINTER :: ifft1_out => NULL()

  COMPLEX(MKC),DIMENSION(:),POINTER :: ifft2_in  => NULL()
  COMPLEX(MKC),DIMENSION(:),POINTER :: ifft2_out => NULL()

  COMPLEX(MKC),DIMENSION(:),POINTER :: ifft3_in  => NULL()
  COMPLEX(MKC),DIMENSION(:),POINTER :: ifft3_out => NULL()

  INTEGER                           :: fft_plan, ifft_plan

  TYPE(class_topology)              :: topo_second

!----------------------------------------------------------------------!
! Initiate
!----------------------------------------------------------------------!
  ierr = 0
  PI = ACOS(-1.0_MK)

  ifft_vort = .FALSE.
  IF(PRESENT(reg_vort) .AND. reg_vort)THEN
    ifft_vort = .TRUE.
  END IF  

!----------------------------------------------------------------------!
! Determine the fft sequence
!----------------------------------------------------------------------!
  fft_seq = patch%fft_seq

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

!----------------------------------------------------------------------!
! Allocate Fourier coefficients for velocity field
!----------------------------------------------------------------------!
  ncell   = topo_second%ncell
  ALLOCATE( vel_fft(pmlib_ndim,ncell(1),ncell(2)) )

!----------------------------------------------------------------------!
! FFT the field in the first two directions and map to the third direction
!----------------------------------------------------------------------!
  fft_seq(2) = 0
  CALL pmlib_fft( topo_all,topo_all%cuboid,mesh%vort,topo_second,vort_fft, &
                & fft_seq,ierr,fft_shift = .FALSE.)

!----------------------------------------------------------------------!
! Setup the third direction
!----------------------------------------------------------------------!
  fft_seq  = patch%fft_seq
  ncell    = topo_second%ncell
  dx       = topo_second%dx
  xmin     = topo_second%xmin

  DO i = 1,pmlib_ndim
    IF( patch%bound_cond(i) .NE. 0 )THEN
      bc_ptc(i) = 1
    ELSE
      bc_ptc(i) = 0
    END IF
  END DO

! Include/exclude ghost cells
  xmin_ptc = patch%xmin - &
           & REAL(( 1 - bc_ptc ) * patch%nghost,MK)*patch%dx
  xmax_ptc = patch%xmax + &
           & REAL(( 1 - bc_ptc ) * patch%nghost,MK)*patch%dx

!----------------------------------------------------------------------!
! Setup wavenumber parameters
!----------------------------------------------------------------------!
  l_ptc(1) = REAL(2 - bc_ptc(1),MK)*(xmax_ptc(1) - xmin_ptc(1))
  l_ptc(2) = REAL(2 - bc_ptc(2),MK)*(xmax_ptc(2) - xmin_ptc(2))

  dk(1) = 2.0_MK*PI/l_ptc(1)
  dk(2) = 2.0_MK*PI/l_ptc(2)

! patch index
  i_ptc = NINT((xmin(1) - xmin_ptc(1))/dx(1))
  j_ptc = NINT((xmin(2) - xmin_ptc(2))/dx(2))

  kmin(1) = - PI/dx(1) + REAL(i_ptc,MK) * dk(1)
  kmin(2) = - PI/dx(2) + REAL(j_ptc,MK) * dk(2)

  eps = pmlib_regularisation_radius*MAXVAL(dx)

!----------------------------------------------------------------------!
! Setup coefficients for the regularisation filter
!----------------------------------------------------------------------!
  c1 = 1.0_MK
  c2 = 0.5_MK
  c3 = 0.125_MK
  c4 = 1.0_MK/48.0_MK
  c5 = 1.0_MK/384.0_MK

!----------------------------------------------------------------------!
! FFT the field in the third direction, convolve and IFFT
!----------------------------------------------------------------------!
  IF(fft_seq(2) .EQ. 1)THEN ! Convolve on x-pencils

    IF(bc_ptc(1) .EQ. 0)THEN
      nfft = 2*ncell(1)
    ELSE
      nfft = ncell(1)
    END IF

! Allocate fft pencils
    ALLOCATE( fft1_in(nfft), fft1_out(nfft), ifft1_in(nfft), ifft1_out(nfft) )
    ALLOCATE( ifft2_in(nfft), ifft2_out(nfft) )
    IF( ifft_vort )THEN
      ALLOCATE( ifft3_in(nfft), ifft3_out(nfft) )
    END IF

    DO j = 1, ncell(2)
! Calculate wave number (FFT-shifted)
      IF( j + j_ptc .LE.  NINT(l_ptc(2)/(2.0_MK*dx(2))) )THEN
        kj = kmin(2) + REAL(j - 1,MK)*dk(2) + PI/dx(2)
      ELSE
        kj = kmin(2) + REAL(j - 1,MK)*dk(2) - PI/dx(2)
      END IF

! Store fft pencil
      DO i = 1, nfft
        IF( i .LE. ncell(1) )THEN
          fft1_in(i) = vort_fft(1,i,j)
        ELSE ! zero padding
          fft1_in(i) = CMPLX(0.0_MK,0.0_MK,MKC)
        END IF
      END DO

! Fourier transform
      CALL dfftw_plan_dft_1d( fft_plan, nfft, fft1_in, fft1_out, & 
                          & FFTW_FORWARD, FFTW_ESTIMATE)
      CALL dfftw_execute_dft( fft_plan, fft1_in, fft1_out )
      CALL dfftw_destroy_plan( fft_plan )

      DO i = 1,nfft
! Calculate wave number in x-direction (FFT-shifted)
        IF( i + i_ptc .LE.  NINT(l_ptc(1)/(2.0_MK*dx(1))) )THEN
          ki = kmin(1) + REAL(i - 1,MK)*dk(1) + PI/dx(1)
        ELSE
          ki = kmin(1) + REAL(i - 1,MK)*dk(1) - PI/dx(1)
        END IF 

! Regularise output the vorticity field
        IF( ifft_vort )THEN
          s = eps*SQRT( ki**2 + kj**2 )

          IF(pmlib_poisson_order .EQ. 0)THEN             
            zeta_fft = CMPLX(1.0_MK,0.0_MK,MKC)
          ELSEIF(pmlib_poisson_order .EQ. 2)THEN
            zeta_fft = CMPLX( c1  &
                     & * EXP(-0.5_MK*s**2) ,0.0_MK,MKC)
          ELSEIF(pmlib_poisson_order .EQ. 4)THEN
            zeta_fft = CMPLX(( c1 + c2*s**2 ) &
                     & * EXP(-0.5_MK*s**2) ,0.0_MK,MKC)
          ELSEIF(pmlib_poisson_order .EQ. 6)THEN
            zeta_fft = CMPLX(( c1 + c2*s**2 + c3*s**4 ) &
                     & * EXP(-0.5_MK*s**2) ,0.0_MK,MKC)
          ELSEIF(pmlib_poisson_order .EQ. 8)THEN
            zeta_fft = CMPLX(( c1 + c2*s**2 + c3*s**4 + c4*s**6 ) &
                     & * EXP(-0.5_MK*s**2) ,0.0_MK,MKC)
          ELSEIF(pmlib_poisson_order .EQ. 10)THEN
            zeta_fft = CMPLX(( c1 + c2*s**2 + c3*s**4 + c4*s**6 &
                     & + c5*s**8 ) * EXP(-0.5_MK*s**2) ,0.0_MK,MKC)
          END IF

! Regularise
          ifft3_in(i) = fft1_out(i) * zeta_fft
        END IF

        IF(pmlib_poisson_kernel .EQ. 1)THEN
! Get the Poisson integration kernel
          G = mesh%pkernel_fft(1,i,j) *dx(1)*dx(2)

! convolve and do spectral curl to get the FFT coefficients of the velocity
          ifft1_in(i) =   G * CMPLX(0.0_MK,kj,MKC) * fft1_out(i)
          ifft2_in(i) = - G * CMPLX(0.0_MK,ki,MKC) * fft1_out(i)

        ELSEIF(pmlib_poisson_kernel .EQ. 2)THEN
! Get the derivative Poisson integration kernel
          Kx = mesh%pkernel_fft(1,i,j) *dx(1)*dx(2)
          Ky = mesh%pkernel_fft(2,i,j) *dx(1)*dx(2)

! Convolve and do spectral curl to get the FFT coefficients of the velocity
          ifft1_in(i) =   Ky * fft1_out(i)
          ifft2_in(i) = - Kx * fft1_out(i)

        END IF

      END DO

! Inverse Fourier transform
      CALL dfftw_plan_dft_1d( ifft_plan, nfft, ifft1_in, ifft1_out, &
                            & FFTW_BACKWARD, FFTW_ESTIMATE)
      CALL dfftw_execute_dft( ifft_plan, ifft1_in, ifft1_out )
      CALL dfftw_execute_dft( ifft_plan, ifft2_in, ifft2_out )
      IF( ifft_vort )THEN
        CALL dfftw_execute_dft( ifft_plan, ifft3_in, ifft3_out )
      END IF
      CALL dfftw_destroy_plan( ifft_plan )

! Store fourier coefficients discard zero-padding
      DO i = 1,ncell(1)
        vel_fft(1,i,j) = ifft1_out(i)/REAL(nfft,MK)
        vel_fft(2,i,j) = ifft2_out(i)/REAL(nfft,MK)
        IF( ifft_vort )THEN
          vort_fft(1,i,j) = ifft3_out(i)/REAL(nfft,MK)
        END IF
      END DO

    END DO !j

! Clean fftw
    CALL dfftw_cleanup()

  ELSEIF(fft_seq(2) .EQ. 2)THEN ! Convolve on y-pencils

    IF(bc_ptc(2) .EQ. 0)THEN
      nfft = 2*ncell(2)
    ELSE
      nfft = ncell(2)
    END IF

! Allocate fft pencils
    ALLOCATE( fft1_in(nfft), fft1_out(nfft), ifft1_in(nfft), ifft1_out(nfft) )
    ALLOCATE( ifft2_in(nfft), ifft2_out(nfft) )
    IF( ifft_vort )THEN
      ALLOCATE( ifft3_in(nfft), ifft3_out(nfft) )
    END IF

    DO i = 1,ncell(1)
! Calculate wave number in x-direction (FFT-shifted)
      IF( i + i_ptc .LE.  NINT(l_ptc(1)/(2.0_MK*dx(1))) )THEN
        ki = kmin(1) + REAL(i - 1,MK)*dk(1) + PI/dx(1)
      ELSE
        ki = kmin(1) + REAL(i - 1,MK)*dk(1) - PI/dx(1)
      END IF 

! Store fft pencil
      DO j = 1, nfft
        IF( j .LE. ncell(2) )THEN
          fft1_in(j) = vort_fft(1,i,j)
        ELSE ! zero padding
          fft1_in(j) = CMPLX(0.0_MK,0.0_MK,MKC)
        END IF
      END DO

! Fourier transform
      CALL dfftw_plan_dft_1d( fft_plan, nfft, fft1_in, fft1_out, & 
                          & FFTW_FORWARD, FFTW_ESTIMATE)
      CALL dfftw_execute_dft( fft_plan, fft1_in, fft1_out )
      CALL dfftw_destroy_plan( fft_plan )

      DO j = 1, nfft
! Calculate wave number in y-direction (FFT-shifted)
        IF( j + j_ptc .LE.  NINT(l_ptc(2)/(2.0_MK*dx(2))) )THEN
          kj = kmin(2) + REAL(j - 1,MK)*dk(2) + PI/dx(2)
        ELSE
          kj = kmin(2) + REAL(j - 1,MK)*dk(2) - PI/dx(2)
        END IF

! Regularise the output vorticity field
        IF( ifft_vort )THEN
          s = eps*SQRT( ki**2 + kj**2 )

          IF(pmlib_poisson_order .EQ. 0)THEN             
            zeta_fft = CMPLX(1.0_MK,0.0_MK,MKC)
          ELSEIF(pmlib_poisson_order .EQ. 2)THEN
            zeta_fft = CMPLX( c1  &
                     & * EXP(-0.5_MK*s**2) ,0.0_MK,MKC)
          ELSEIF(pmlib_poisson_order .EQ. 4)THEN
            zeta_fft = CMPLX(( c1 + c2*s**2 ) &
                     & * EXP(-0.5_MK*s**2) ,0.0_MK,MKC)
          ELSEIF(pmlib_poisson_order .EQ. 6)THEN
            zeta_fft = CMPLX(( c1 + c2*s**2 + c3*s**4 ) &
                     & * EXP(-0.5_MK*s**2) ,0.0_MK,MKC)
          ELSEIF(pmlib_poisson_order .EQ. 8)THEN
            zeta_fft = CMPLX(( c1 + c2*s**2 + c3*s**4 + c4*s**6 ) &
                     & * EXP(-0.5_MK*s**2) ,0.0_MK,MKC)
          ELSEIF(pmlib_poisson_order .EQ. 10)THEN
            zeta_fft = CMPLX(( c1 + c2*s**2 + c3*s**4 + c4*s**6 &
                     & + c5*s**8 ) * EXP(-0.5_MK*s**2) ,0.0_MK,MKC)
          END IF

! Regularise
          ifft3_in(j) = fft1_out(j) * zeta_fft
        END IF

        IF(pmlib_poisson_kernel .EQ. 1)THEN
! Get the Poisson integration kernel
          G = mesh%pkernel_fft(1,i,j) *dx(1)*dx(2)

! Convolve and do spectral curl to get the FFT coefficients of the velocity
          ifft1_in(j) =   G * CMPLX(0.0_MK,kj,MKC) * fft1_out(j)
          ifft2_in(j) = - G * CMPLX(0.0_MK,ki,MKC) * fft1_out(j) 

        ELSEIF(pmlib_poisson_kernel .EQ. 2)THEN
! Get the derivative Poisson integration kernel
          Kx = mesh%pkernel_fft(1,i,j) *dx(1)*dx(2)
          Ky = mesh%pkernel_fft(2,i,j) *dx(1)*dx(2)

! Convolve and do spectral curl to get the FFT coefficients of the velocity
          ifft1_in(j) =   Ky * fft1_out(j)
          ifft2_in(j) = - Kx * fft1_out(j)

        END IF
      END DO

! Inverse Fourier transform
      CALL dfftw_plan_dft_1d( ifft_plan, nfft, ifft1_in, ifft1_out, &
                            & FFTW_BACKWARD, FFTW_ESTIMATE)
      CALL dfftw_execute_dft( ifft_plan, ifft1_in, ifft1_out )
      CALL dfftw_execute_dft( ifft_plan, ifft2_in, ifft2_out )
      IF( ifft_vort )THEN
        CALL dfftw_execute_dft( ifft_plan, ifft3_in, ifft3_out )
      END IF
      CALL dfftw_destroy_plan( ifft_plan )

! Store fourier coefficients discard zero-padding
      DO j = 1,ncell(2)
        vel_fft(1,i,j) = ifft1_out(j)/REAL(nfft,MK)
        vel_fft(2,i,j) = ifft2_out(j)/REAL(nfft,MK)
        IF( ifft_vort )THEN
          vort_fft(1,i,j) = ifft3_out(j)/REAL(nfft,MK)
        END IF
      END DO

    END DO !i

! Clean fftw
    CALL dfftw_cleanup()

  END IF

!----------------------------------------------------------------------!
! De-allocate fft arrays
!----------------------------------------------------------------------!
  DEALLOCATE( fft1_in, fft1_out )
  DEALLOCATE( ifft1_in, ifft1_out )
  DEALLOCATE( ifft2_in, ifft2_out )

  IF( ifft_vort )THEN
    DEALLOCATE( ifft3_in, ifft3_out )
  END IF

!----------------------------------------------------------------------!
! IFFT velocity field in the first direction
!----------------------------------------------------------------------!
  ifft_seq(1) = patch%fft_seq(1)
  ifft_seq(2) = 0
  CALL pmlib_ifft( topo_all,topo_second,vel_fft,topo_all%cuboid,mesh%vel, &
                 & ifft_seq,ierr,fft_shift = .FALSE.)

  IF( ifft_vort )THEN
    CALL pmlib_ifft( topo_all,topo_second,vort_fft,topo_all%cuboid,mesh%vort, &
                   & ifft_seq,ierr,fft_shift = .FALSE.)
  END IF

!----------------------------------------------------------------------!
! Deallocate local pointers
!----------------------------------------------------------------------!
  IF( ASSOCIATED(vort_fft) )THEN 
    DEALLOCATE(vort_fft)
  END IF
  IF( ASSOCIATED(vel_fft) ) THEN
    DEALLOCATE(vel_fft)
  END IF

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