module wei_fft
  !! FFTW
  !! used for data(mx,my,hx,hy,w)
#ifndef FFTW
use fftw_sep_fake_mod
#endif


  use wei_process
  use wei_flags
  use wei_attr
  !---------------
  implicit none
   integer, parameter :: max_p=200
  integer(kind=dbl), private,save :: planF(max_p) !! forward FFT
  integer(kind=dbl), private,save :: planI(max_p) !! inverse FFT
  integer, private,save ::  nplans,ndim_p(5,max_p),cur_plan
  real,      private :: scale
  logical,private,save :: fftmeasure

  !! temp  arrays
  complex, dimension(:,:), allocatable, private :: mmhht
  complex, dimension(:,:,:),         pointer, private ::   hht

  !! shift arrays
  complex, dimension(:,:),           pointer, private :: sftF
  complex, dimension(:,:),           pointer, private :: sftI

contains
  !----------------------------------------------------------------
  subroutine weifft_init()
    integer, dimension(5) :: ni
    real    :: arg
    integer :: ihx,ihy
    integer :: er
#ifdef DEBUG
    character(len=128), parameter :: unit='weifft_init'
    if(pro%debug) call in(unit)
#endif

    call from_param("fftmeasure",fftmeasure,.false.)

    !! FFT
    ni = (/ pro%amx%n, pro%amy%n, pro%ahx%n, pro%ahy%n, 1 /)
    scale = 1./sqrt(1.*product(ni))

    allocate(mmhht(ni(1)*ni(2)*ni(3)*ni(4),pro%nth),stat=er)
    if(er/=0) call seperr('cant allocate mmhht')

    nplans=0
    cur_plan=0

#ifdef DEBUG
    if(pro%debug) call out(unit)
#endif
  end subroutine weifft_init

 integer function find_plan(ni)
   integer  :: i,ni(5)
   logical :: found

    i=1;
   found=.false.
    do while(i <= nplans .and. .not. found) 
      if(all(ni == ndim_p(:,i))) then
       found=.true.
      else
       i=i+1
     end if
    end do
     if(.not. found) then
       call add_plan(ni)
       i=cur_plan
     end if
     find_plan=i
 end function find_plan

 subroutine add_plan(ni)
  integer, dimension(5) :: ni
    if(nplans<max_p) then
      cur_plan=cur_plan+1
      nplans=nplans+1
    else 
     if(cur_plan==max_p) then
       cur_plan=1
     else
       cur_plan=cur_plan+1
     end if 
      call fftwnd_f77_destroy_plan(planF(cur_plan))
      call fftwnd_f77_destroy_plan(planI(cur_plan))
    end if

    ndim_p(:,cur_plan)=ni

    if(fftmeasure) then
       !! init  forward plan
       call fftwnd_f77_create_plan(planF(cur_plan),5,ni,fFTW_FORWARD  &
       ,    fFTW_MEASURE + fFTW_THREADSAFE + fFTW_IN_PLACE)
       !! init backward plan
       call fftwnd_f77_create_plan(planI(cur_plan),5,ni,fFTW_BACKWARD &
       ,    fFTW_MEASURE + fFTW_THREADSAFE + fFTW_IN_PLACE) 
    else
       !! init  forward plan
       call fftwnd_f77_create_plan(planF(cur_plan),5,ni,fFTW_FORWARD  &
       ,    fFTW_ESTIMATE + fFTW_THREADSAFE + fFTW_IN_PLACE)
       !! init backward plan
       call fftwnd_f77_create_plan(planI(cur_plan),5,ni,fFTW_BACKWARD &
       ,    fFTW_ESTIMATE + fFTW_THREADSAFE + fFTW_IN_PLACE)
    end if
  end subroutine add_plan 



  !----------------------------------------------------------------
  !! forward FFT
  subroutine weifft(iws,izs,ith,data)
    integer, intent(in) :: iws,ith,izs
    complex, dimension(:,:,:,:,:), pointer :: data
    integer :: imx,imy
    integer :: iplan
    integer :: nd(4)
#ifdef DEBUG
    character(len=128), parameter :: unit='weifft'
    if(pro%depar) call in(unit)
#endif

  

     nd=(/size(data,1),size(data,2), pro%ihxb_end(izs) -pro%ihxb_beg(izs)+1, &
     pro%ihyb_end(izs) -pro%ihyb_beg(izs)+1/)


      iplan=find_plan((/nd(1),nd(2),nd(3),nd(4),1/))

    !! FFT
    mmhht(1:product(nd),ith) =  reshape( &
    data(:,:,pro%ihxb_beg(izs):pro%ihxb_end(izs),pro%ihyb_beg(izs):pro%ihyb_end(izs),iws),(/product(nd)/))

    call fftwnd_f77_one(planF(iplan),mmhht(:,ith),0)

    data(:,:,pro%ihxb_beg(izs):pro%ihxb_end(izs),pro%ihyb_beg(izs):pro%ihyb_end(izs),iws)=&
    reshape(mmhht(1:product(nd),ith),nd)* 1./sqrt(1.*product(nd))


#ifdef DEBUG
    if(pro%depar) call out(unit)
#endif
  end subroutine weifft
  !----------------------------------------------------------------
  !! inverse FFT
  subroutine weiift(iws,izs,ith,data)
    integer, intent(in) :: iws, ith,izs
    complex, dimension(:,:,:,:,:), pointer :: data
    integer :: imx,imy,iplan,nd(4)
#ifdef DEBUG
    character(len=128), parameter :: unit='weiift'
    if(pro%depar) call in(unit)
#endif

     nd=(/size(data,1),size(data,2), pro%ihxb_end(izs) -pro%ihxb_beg(izs)+1, &
     pro%ihyb_end(izs) -pro%ihyb_beg(izs)+1/)

      iplan=find_plan((/nd(1),nd(2),nd(3),nd(4),1/))

    !! FFT
    mmhht(1:product(nd),ith) =  reshape( &
    data(:,:,pro%ihxb_beg(izs):pro%ihxb_end(izs),pro%ihyb_beg(izs):pro%ihyb_end(izs),iws),(/product(nd)/))

    call fftwnd_f77_one(planI(iplan),mmhht(:,ith),0) !! Xform in place

    data(:,:,pro%ihxb_beg(izs):pro%ihxb_end(izs),pro%ihyb_beg(izs):pro%ihyb_end(izs),iws)=&
    reshape(mmhht(1:product(nd),ith),nd)* 1./sqrt(1.*product(nd))

#ifdef DEBUG
    if(pro%depar) call out(unit)
#endif
  end subroutine weiift
  !----------------------------------------------------------------
  !! forward FFT
  subroutine weifft4d(data)
    complex, dimension(:,:,:,:), pointer :: data
#ifdef DEBUG
    character(len=128), parameter :: unit='weifft4d'
    if(pro%depar) call in(unit)
#endif

    !! FFT
    mmhht(:,1) =reshape( data,(/size(mmhht)/))
    call fftwnd_f77_one(planF(1),mmhht(:,1),0)
    data(:,:,:,:) =reshape( mmhht(:,1) * 1./sqrt(1.*size(data)),(/size(data,1),&
      size(data,2),size(data,3),size(data,4)/))

#ifdef DEBUG
    if(pro%depar) call out(unit)
#endif
  end subroutine weifft4d
  !----------------------------------------------------------------
  !! inverse FFT
  subroutine weiift4d(data)
    complex, dimension(:,:,:,:), pointer :: data
#ifdef DEBUG
    character(len=128), parameter :: unit='weiift4d'
    if(pro%depar) call in(unit)
#endif

    !! IFT
    mmhht(:,1) =reshape( data,(/size(mmhht)/))
    call fftwnd_f77_one(planI(1),mmhht(:,1),0) !! Xform in place
    data(:,:,:,:) =reshape( mmhht(:,1) * 1./sqrt(1.*size(data)),(/size(data,1),&
      size(data,2),size(data,3),size(data,4)/))

#ifdef DEBUG
    if(pro%depar) call out(unit)
#endif
  end subroutine weiift4d
  !----------------------------------------------------------------
end module wei_fft
