!=======================================================================
!
!  PROGRAM  PHASE/0 2018.01 ($Rev: 570 $)
!
!  MODULE: m_FFT
!
!  AUTHOR(S): T. Yamasaki, K. Betsuyaku,   August/20/2003
!  
!  FURTHER MODIFICATION: T. Yamasaki, January/13/2004
!                                   , March/14/2005, April/10/2007
!
!  Contact address :  Phase System Consortium
!                     E-mail: phase_system@nims.go.jp URL https://azuma.nims.go.jp
!  
!
!
!=======================================================================
!
!     The original version of this set of the computer programs "PHASE"
!  was developed by the members of the Theory Group of Joint Research
!  Center for Atom Technology (JRCAT), based in Tsukuba, in the period
!  1993-2001.
!
!     Since 2002, this set has been tuned and new functions have been
!  added to it as a part of the national project "Frontier Simulation 
!  Software for Industrial Science (FSIS)",  which is supported by
!  the IT program of the Ministry of Education, Culture, Sports,
!  Science and Technology (MEXT) of Japan. 
!     Since 2006, this program set has been developed as a part of the
!  national project "Revolutionary Simulation Software (RSS21)", which
!  is supported by the next-generation IT program of MEXT of Japan.
!   Since 2013, this program set has been further developed centering on PHASE System
!  Consortium.
!   The activity of development of this program set has been supervised by Takahisa Ohno.
!
!
  integer                   :: idp, nlp, nmp, nnp

  complex(kind=CMPLDP),private,target,allocatable,dimension(:)    :: cw1,cw2,cw3
  complex(kind=CMPLDP),private,allocatable,dimension(:)    :: wlp,wmp,wnp
  integer,             private,allocatable,dimension(:,:)  :: iwork

  ! ------- Positron start 
  complex(kind=CMPLDP),private,target,allocatable,dimension(:)    :: cw1_pstrn,cw2_pstrn,cw3_pstrn

  integer :: sw_avoiding_odd_fftbox = OFF
  integer :: sw_zero_padding = OFF

  real(kind=DP),private,allocatable,dimension(:)           :: ftw

  include 'mpif.h'

contains
  subroutine fft_WFCD_work_alloc
    integer :: nid,nnl,nnm,nnn
    integer :: istat = 0
! ----> work arrays of fft for wave functions
    nnl  = fft_box_size_WF(1,1)
    nnm  = fft_box_size_WF(2,1);  nnn  = fft_box_size_WF(3,1)
    if(ipri >= 1) write(nfout,'(" !!allocation of cw1,cw2,cw3 (fft_WFCD_work_alloc)")')
    allocate(cw1(nnl), stat=istat)
    allocate(cw2(nnm), stat=istat)
    allocate(cw3(nnn), stat=istat)
    if(ipri >= 1) write(nfout,'(" !! WF_JRCATFFT <<m_FFT.fft_WFCD_work_alloc>>")')

! ------- Positron start 
! ----> and work arrays of fft for positron wave functions
    if(sw_positron /= OFF) then
       nnl  = fft_box_size_pWF(1,1)
       nnm  = fft_box_size_pWF(2,1);  nnn  = fft_box_size_pWF(3,1)
       if(ipri >= 1) write(nfout,'(" !!allocation of cw1_pstrn,cw2_pstrn,cw3_pstrn (fft_WFCD_work_alloc)")')
       allocate(cw1_pstrn(nnl), stat=istat)
       allocate(cw2_pstrn(nnm), stat=istat)
       allocate(cw3_pstrn(nnn), stat=istat)
    end if
! ------- Positron end

! ----> work arrays of ifft for the charge density
    nid = fft_box_size_CD(1,0);  nnl = fft_box_size_CD(1,1)
    nnm = fft_box_size_CD(2,1);  nnn = fft_box_size_CD(3,1)
    allocate(wlp(3*nnl+7), stat=istat)
    allocate(wmp(2*nnm*(nid+1)+7), stat=istat)
    allocate(wnp(2*nnn+7), stat=istat)
    allocate(iwork(2,nnl+nnm+nnn), stat=istat)
  end subroutine fft_WFCD_work_alloc

  subroutine m_FFT_alloc_WF_work
    integer :: istat = 0
    allocate(ftw(nfft), stat=istat)
  end subroutine m_FFT_alloc_WF_work

  subroutine m_FFT_alloc_pWF_work()
    integer :: istat = 0
    allocate(ftw(nfft_pstrn), stat=istat)
  end subroutine m_FFT_alloc_pWF_work

  subroutine m_FFT_dealloc_WF_work
    integer :: istat = 0
    if(allocated(ftw)) then
       deallocate(ftw, stat=istat)
       if(istat /= 0 ) then
          if(ipri>=1) then
             write(nfout,*) 'Deallocation error for ftw in sub. m_FFT_dealloc_WF_work'
             write(nfout,*) 'stat=', istat
          end if
          stop
       end if
    end if
  end subroutine m_FFT_dealloc_WF_work

  subroutine m_FFT_alloc_CD_box
    integer :: istat = 0
    allocate(afft_CD(nfftp), stat=istat)
  end subroutine m_FFT_alloc_CD_box

  subroutine m_FFT_dealloc_CD_box
    integer :: istat = 0
    if(allocated(ftw)) then
       deallocate(ftw, stat=istat)
    end if
    deallocate(afft_CD, stat=istat)

  end subroutine m_FFT_dealloc_CD_box

  subroutine m_FFT_setup(inversion_symmetry,paramset)
    integer, intent(in) :: inversion_symmetry
    logical, intent(in) :: paramset

    real(kind=DP),allocatable,dimension(:)               :: cfft
    integer :: nfft_t, ipad
    integer :: istat = 0
    integer :: id_sname = -1
    call tstatc0_begin('m_FFT_setup ',id_sname)

    if(inversion_symmetry == ON) then  ! kimg == 1
       ipad = 2
    else if(inversion_symmetry == OFF) then ! kimg == 2
       ipad = 1
    end if
! --- fft_box_size_WF, fft_box_size_pWF ---
    fft_box_size_WF(1,0) = fft_box_size_WF(1,1) + ipad
    fft_box_size_WF(2:3,0) = fft_box_size_WF(2:3,1)
    if(sw_positron /= OFF) then
       fft_box_size_pWF(1,0)   = fft_box_size_pWF(1,1) + ipad
       fft_box_size_pWF(2:3,0) = fft_box_size_pWF(2:3,1)
    end if

    nfft =   product(fft_box_size_WF(1:3,0)) * (2-inversion_symmetry)
! --- fft_box_size_CD ---
    fft_box_size_CD_nonpara(1,0)   = fft_box_size_CD(1,1) + ipad
    fft_box_size_CD_nonpara(2:3,0) = fft_box_size_CD(2:3,1)
    nfftp_nonpara  = product(fft_box_size_CD_nonpara(1:3,0)) * (2-inversion_symmetry)

    fft_box_size_CD(1:3,0) = fft_box_size_CD_nonpara(1:3,0)
    fft_box_size_CD_c(1:3,0) = fft_box_size_CD(1:3,0)
    nfftp  = product(fft_box_size_CD(1:3,0)) * (2-inversion_symmetry)
    nfftps = nfftp

    if(ipri >= 1) call wd_FFTboxsizes(nfout)

    idp = fft_box_size_CD(1,0)
    nlp = fft_box_size_CD(1,1)
    nmp = fft_box_size_CD(2,1)
    nnp = fft_box_size_CD(3,1)

    if(sw_positron /= OFF) &
         & nfft_pstrn = product(fft_box_size_pWF(1:3,0))*(2-inversion_symmetry)

    if(.not. paramset) then
       call fft_WFCD_work_alloc     ! <cw[123],cw[123]_pstrn,wlp,wmp,wnp> are allocated

! Initialization of the Wave-Function FFT
       if(sw_positron /= OFF .and. nfft_pstrn > nfft) then
          call m_FFT_alloc_pWF_work()  ! <ftw> is allocated
          nfft_t = nfft_pstrn
       else
          call m_FFT_alloc_WF_work()   ! <ftw> is allocated
          nfft_t = nfft
       end if
       allocate(cfft(nfft_t), stat=istat) ! cfft is used only for initiallization
       if(istat /= 0) then
          if(ipri >= 1) then
             write(nfout,*) 'Allocation error for cfft in sub. m_FFT_setup'
             write(nfout,*) 'stat =', istat, 'nfft =', nfft_t
          end if
          stop
       end if

       call init_fft_coefficients_arrays_WF()

       call m_FFT_dealloc_WF_work()
       deallocate(cfft,stat=istat)
       if(istat /= 0 ) then
          if(ipri >= 1) then
             write(nfout,*) 'Deallocation error for cfft in sub. m_FFT_setup'
             write(nfout,*) 'stat =', istat
          end if
          stop
       end if

       call CDFFT_setup()

    endif

    call tstatc0_end(id_sname)
  contains
    subroutine init_fft_coefficients_arrays_WF()
      integer :: id, nl, nm, nn, ierr
      integer :: id_p, nl_p, nm_p, nn_p
      id = fft_box_size_WF(1,0)
      nl = fft_box_size_WF(1,1)
      nm = fft_box_size_WF(2,1)
      nn = fft_box_size_WF(3,1)
      if(sw_positron /= OFF) then
         id_p = fft_box_size_pWF(1,0)
         nl_p = fft_box_size_pWF(1,1)
         nm_p = fft_box_size_pWF(2,1)
         nn_p = fft_box_size_pWF(3,1)
      end if

      if(kimg == 1) then
         call jrcat_r3ft(cfft,ftw,id,nl,nm,nn,cw1,cw2,cw3,0,0,0,0)
         if(sw_positron /= OFF) &
              & call jrcat_r3ft(cfft,ftw,id_p,nl_p,nm_p,nn_p &
              &                   ,cw1_pstrn,cw2_pstrn,cw3_pstrn,0,0,0,0)
      else  ! kimg == 2
         call jrcat_c3ft(cfft,ftw,id,nl,nm,nn,cw1,cw2,cw3,0,0,0,0)
         if(ipri >= 2) then
            write(nfout,'(" !! cw1,cw2 and cw3 <<m_FFT.init_fft_coefficients_arrays_WF>>")')
            write(nfout,'(" !! cw1 <<m_FFT.init_fft_coefficients_arrays_WF>>")')
            write(nfout,'(8f8.4)') cw1
            write(nfout,'(" !! cw2 <<m_FFT.init_fft_coefficients_arrays_WF>>")')
            write(nfout,'(8f8.4)') cw2
            write(nfout,'(" !! cw3 <<m_FFT.init_fft_coefficients_arrays_WF>>")')
            write(nfout,'(8f8.4)') cw3
         end if
         if(sw_positron /= OFF) call jrcat_c3ft(cfft,ftw,id_p,nl_p,nm_p,nn_p &
              &                                 ,cw1_pstrn,cw2_pstrn,cw3_pstrn,0,0,0,0)
      endif
    end subroutine init_fft_coefficients_arrays_WF
  end subroutine m_FFT_setup

  subroutine CDFFT_setup()
    call m_FFT_alloc_CD_box()       ! <ftw> is allocated

    call init_fft_coefficients_arrays_CD()

    CD_setup_is_done = YES

    call m_FFT_dealloc_CD_box()
  contains
    subroutine init_fft_coefficients_arrays_CD
      integer :: ierr
   !  ---> FFT for Charge density
      if(kimg == 1) then
         call r3fft(afft_CD,idp,nlp,nmp,nnp,wlp,wmp,wnp,0,0,1,ftw,ierr)
      else
         call c3fft(afft_CD,idp,nlp,nmp,nnp,wlp,wmp,wnp,0,0,1,ftw,ierr)
      endif
    end subroutine init_fft_coefficients_arrays_CD
  end subroutine CDFFT_setup

  subroutine m_FFT_WF(electron_or_positron,nfout,afft,inverse_or_direct,switch)  ! G space --> R space
    integer, intent(in)          :: electron_or_positron
    integer, intent(in)          :: nfout
    real(kind=DP), intent(inout) :: afft(nfft)
    integer, intent(in)          :: inverse_or_direct
    integer, intent(in)          :: switch

    complex(kind=CMPLDP),pointer,dimension(:) :: cw1_t, cw2_t, cw3_t

    integer :: id, nl, nm, nn

    integer, dimension(2) :: flag_jrcat_r3ft = (/-2,  2/)
    integer, dimension(2) :: flag_jrcat_c3ft = (/ 1, -1/)
    integer kc1, kc2, kc3

    integer :: id_sname = -1
    if(electron_or_positron == ELECTRON) then
       call tstatc0_begin('m_FFT_WF ',id_sname)
       id = fft_box_size_WF(1,0)
       nl = fft_box_size_WF(1,1)
       nm = fft_box_size_WF(2,1)
       nn = fft_box_size_WF(3,1)
       cw1_t=>cw1; cw2_t=>cw2; cw3_t=>cw3 
    else if(electron_or_positron == POSITRON) then
       call tstatc0_begin('m_FFT_pWF ',id_sname)
       id = fft_box_size_pWF(1,0)
       nl = fft_box_size_pWF(1,1)
       nm = fft_box_size_pWF(2,1)
       nn = fft_box_size_pWF(3,1)
       cw1_t=>cw1_pstrn; cw2_t=>cw2_pstrn; cw3_t=>cw3_pstrn
    end if

    if(switch == ON) then
       kc1 = 0;       kc2 = nm/4 + 1;       kc3 = nn/4 + 1
    else
       kc1 = 0;       kc2 = 0       ;       kc3 = 0
    endif
    if(kimg == 1) then
       call jrcat_r3ft(afft,ftw,id,nl,nm,nn,cw1_t,cw2_t,cw3_t,kc1,kc2,kc3&
            & ,flag_jrcat_r3ft(inverse_or_direct))
    else if(kimg == 2) then
       call jrcat_c3ft(afft,ftw,id,nl,nm,nn,cw1_t,cw2_t,cw3_t,kc1,kc2,kc3&
            & ,flag_jrcat_c3ft(inverse_or_direct))
    endif

    call tstatc0_end(id_sname)

  end subroutine m_FFT_WF

  subroutine fft_CD_inverse_core(afft_CD)
    real(kind=DP),intent(inout),dimension(nfftp_nonpara) :: afft_CD
    integer :: iopt = 0
    integer :: ier1 = 0

    if(kimg == 1) then
       call r3fft(afft_CD,idp,nlp,nmp,nnp,wlp,wmp,wnp,iopt,-1,1,iwork,ier1)
    else if(kimg == 2) then
       call c3fft(afft_CD,idp,nlp,nmp,nnp,wlp,wmp,wnp,iopt, 1,1,iwork,ier1)
    end if

    if (ier1 /= 0) then
       if(ipri>=1) write (nfout,*) ' !!Error fft(inverse) ier1 = ',ier1
       stop
    end if

  end subroutine fft_CD_inverse_core

  subroutine fft_CD_direct_core(afft_CD)
    real(kind=DP),intent(inout),dimension(nfftp_nonpara) :: afft_CD
    integer :: iopt = 0
    integer :: ier1 = 0

    if(kimg == 1) then
       call r3fft(afft_CD,idp,nlp,nmp,nnp,wlp,wmp,wnp,iopt,1,1,iwork,ier1)
    else if(kimg == 2) then
       call c3fft(afft_CD,idp,nlp,nmp,nnp,wlp,wmp,wnp,iopt,-1,1,iwork,ier1)
    end if

    if (ier1 /= 0) then
       if(ipri>=1) write (nfout,*) ' !!Error  fft(direct) ier1 = ',ier1
       stop
    end if
  end subroutine fft_CD_direct_core
