module fft_data_mod

  implicit none

  ! GOOD For LINUX
  integer, parameter ::  FFT_FORW=-1
  integer, parameter ::  FFT_BACK=1

  ! Parameters for FFTW package
  integer FFTW_FORWARD,FFTW_BACKWARD
  parameter (FFTW_FORWARD=-1,FFTW_BACKWARD=1)

  integer FFTW_REAL_TO_COMPLEX,FFTW_COMPLEX_TO_REAL
  parameter (FFTW_REAL_TO_COMPLEX=-1,FFTW_COMPLEX_TO_REAL=1)

  integer FFTW_ESTIMATE,FFTW_MEASURE
  parameter (FFTW_ESTIMATE=0,FFTW_MEASURE=1)

  integer FFTW_OUT_OF_PLACE,FFTW_IN_PLACE,FFTW_USE_WISDOM
  parameter (FFTW_OUT_OF_PLACE=0)
  parameter (FFTW_IN_PLACE=8,FFTW_USE_WISDOM=16)

  integer FFTW_THREADSAFE
  parameter (FFTW_THREADSAFE=128)

  type fft_type 
     integer   :: ierr,ndim
!!!#if defined (LINUX)
!     integer   :: plan_forw,plan_back
!!!!#else
     integer*8 :: plan_forw,plan_back
!!!#endif
     real :: scale
  end type fft_type

  integer, external :: omp_get_thread_num

contains

  ! ----------------------------------------------------------------------
  function  init_fft_1d(n1,fast_init) result (fft_id)
    type (fft_type) fft_id
    integer n1,stat,flag
    logical, optional :: fast_init

    if(present(fast_init)) then
       if(fast_init) then
          flag=ior(ior(FFTW_ESTIMATE,FFTW_USE_WISDOM),FFTW_THREADSAFE)
       else
          flag=ior(ior(FFTW_MEASURE,FFTW_USE_WISDOM),FFTW_THREADSAFE)
       end if
    else
       flag=ior(ior(FFTW_ESTIMATE,FFTW_USE_WISDOM),FFTW_THREADSAFE)
    end if
    fft_id%scale=sqrt(1./(n1))
    call fftw_plan_dft_1d(fft_id%plan_forw,n1,FFTW_FORWARD ,flag)
    call fftw_plan_dft_1d(fft_id%plan_back,n1,FFTW_BACKWARD,flag)
    if((fft_id%plan_forw == 0 ).or.(fft_id%plan_back == 0)) then
       fft_id%ierr=1
    else
       fft_id%ierr=0
    end if
    fft_id%ndim=1

  end function init_fft_1d
  ! ----------------------------------------------------------------------
!  function init_rfft_1d(n1) result (fft_id)
!    type(fft_type)     :: fft_id
!    integer, intent(in):: n1
!    integer            :: stat, flag
!
!    flag=ior(ior(FFTW_ESTIMATE,FFTW_USE_WISDOM),FFTW_THREADSAFE)
!    fft_id%scale=sqrt(1./(n1))
!    call rfftw_f77_create_plan(fft_id%plan_forw,n1,FFTW_REAL_TO_COMPLEX,flag)
!    call rfftw_f77_create_plan(fft_id%plan_back,n1,FFTW_COMPLEX_TO_REAL,flag)
!    if((fft_id%plan_forw == 0 ).or.(fft_id%plan_back == 0)) then
!       fft_id%ierr=1
!    else
!       fft_id%ierr=0
!    end if
!    fft_id%ndim=1
!  end function init_rfft_1d
  ! ----------------------------------------------------------------------
  function  init_fft_2d(n1,n2,fast_init) result (fft_id)
    type (fft_type) fft_id
    integer n1,n2,stat,flag
    logical, optional :: fast_init

    if(present(fast_init)) then
       if(fast_init) then
          flag=ior(ior(FFTW_ESTIMATE,FFTW_USE_WISDOM),FFTW_THREADSAFE)
       else
          flag=ior(ior(FFTW_MEASURE,FFTW_USE_WISDOM),FFTW_THREADSAFE)
       end if
    else
       flag=ior(ior(FFTW_ESTIMATE,FFTW_USE_WISDOM),FFTW_THREADSAFE)
    end if

    fft_id%scale=sqrt(1./(n1*n2))
    call fftw_plan_dft_2d(fft_id%plan_forw,n1,n2,FFTW_FORWARD ,flag)
    call fftw_plan_dft_2d(fft_id%plan_back,n1,n2,FFTW_BACKWARD,flag)

    if((fft_id%plan_forw ==0 ).or.(fft_id%plan_back == 0)) then
       fft_id%ierr=1
    else
       fft_id%ierr=0
    end if

    fft_id%ndim=2
  end function init_fft_2d

  !----------------------------------------------------------------------
  function init_fft_3d(n1,n2,n3,fast_init) result (fft_id)
    type (fft_type) fft_id
    integer n1,n2,n3,stat,flag
    logical, optional :: fast_init


    if(present(fast_init)) then
       if(fast_init) then
          flag=ior(ior(FFTW_ESTIMATE,FFTW_USE_WISDOM),FFTW_THREADSAFE)
       else
          flag=ior(ior(FFTW_MEASURE,FFTW_USE_WISDOM),FFTW_THREADSAFE)
       end if
    else
       flag=ior(ior(FFTW_ESTIMATE,FFTW_USE_WISDOM),FFTW_THREADSAFE)
    end if

    fft_id%scale=sqrt(1./(n1*n2*n3))

    call fftw_plan_dft_3d(fft_id%plan_forw,n1,n2,n3,FFTW_FORWARD,flag)
    call fftw_plan_dft_3d(fft_id%plan_back,n1,n2,n3,FFTW_BACKWARD,flag)

    if((fft_id%plan_forw ==0 ).or.(fft_id%plan_back == 0)) then
       fft_id%ierr=1
    else
       fft_id%ierr=0
    end if

    fft_id%ndim=3
  end function init_fft_3d
  !----------------------------------------------------------------------
  subroutine fft_1d_data(fft_id,sign,n1,data_vec)
    type (fft_type) fft_id
    integer :: sign, n1
    complex, dimension (:) :: data_vec
    complex, dimension (n1) :: temp_1d
 
    !    if(omp_get_thread_num() == 0) write(0,*)'Beg fft 1d'
    temp_1d=data_vec*fft_id%scale
    if(sign == FFT_FORW) then
       call fftw_f77_one(fft_id%plan_forw, temp_1d,data_vec)
    else
       !       if(omp_get_thread_num() == 0) write(0,*)'Back 1d'
       call fftw_f77_one(fft_id%plan_back, temp_1d,data_vec)
    end if
    return	
  end subroutine fft_1d_data
  !----------------------------------------------------------------------
  subroutine rfft_1d_data(fft_id,sign,n,real_vec,cmpx_vec)
  ! No side effects. cmpx_vec has both zero freq and Nyq
 
    type(fft_type),            intent(in   ) :: fft_id
    integer,                   intent(in   ) :: sign, n
    real,    dimension(n    )                :: real_vec
    complex, dimension(n/2+1), intent(inout) :: cmpx_vec
    real,    dimension(n    )                :: temp_vec
    integer                                  :: i

    if(mod(n,2)==1) then
       write(0,*) "Error! rfft_1d_data needs a even n. Now, n =",n
    end if
    
    if(sign == FFTW_REAL_TO_COMPLEX) then
       call rfftw_f77_one(fft_id%plan_forw, real_vec, temp_vec)
       cmpx_vec(    1)=cmplx(temp_vec(    1),0.) ! zero freq
       cmpx_vec(n/2+1)=cmplx(temp_vec(n/2+1),0.) ! Nyquist
       do i=2,n/2
             cmpx_vec(i)=cmplx(temp_vec(i),-temp_vec(n+2-i))
       end do
       cmpx_vec=cmpx_vec*fft_id%scale
    else
       temp_vec(:n/2+1)=real(cmpx_vec)
       do i=0,n/2-2
          temp_vec(n-i)=-aimag(cmpx_vec(i+2))
       end do
       call rfftw_f77_one(fft_id%plan_back, temp_vec, real_vec)
       real_vec=real_vec*fft_id%scale
    end if
    return	
  end subroutine rfft_1d_data 
  !----------------------------------------------------------------------
  subroutine fft_2d_data(fft_id,sign,n1,n2,data_slice)
    type (fft_type) fft_id
    integer :: n1,n2
    integer :: sign
    complex, dimension (:,:) :: data_slice
    complex, dimension (n1,n2) :: temp_2d

    temp_2d=data_slice*fft_id%scale
    if(sign == FFT_FORW) then
       call fftwnd_f77_one(fft_id%plan_forw, temp_2d,data_slice)
    else
       call fftwnd_f77_one(fft_id%plan_back, temp_2d,data_slice)
    end if

    return	
  end subroutine fft_2d_data
  !----------------------------------------------------------------------
  subroutine fft_3d_data(fft_id,sign,n1,n2,n3,data_slice)
    type (fft_type) fft_id
    integer :: n1,n2,n3
    integer :: sign
    complex,  dimension (:,:,:) :: data_slice
    complex,  dimension (n1,n2,n3) :: temp_3d

    temp_3d=data_slice*fft_id%scale

    if(sign == FFT_FORW) then
       call fftwnd_f77_one(fft_id%plan_forw, temp_3d,data_slice)
    else
       call fftwnd_f77_one(fft_id%plan_back, temp_3d,data_slice)
    end if

    return	
  end subroutine fft_3d_data

  !----------------------------------------------------------------------
  subroutine fft_3d_two_pass_data(fft_1d_id,fft_2d_id,sign,n1,n2,n3,data_slice)
    type (fft_type) fft_1d_id,fft_2d_id
    integer :: n1,n2,n3,i1,i2,i3
    integer :: sign
    real    :: scale
    complex,  dimension (:,:,:) :: data_slice
    complex,  dimension (n1,n2) :: temp_2d
    complex,  dimension (n3)    :: temp_1d

    scale=fft_2d_id%scale*fft_1d_id%scale

    !temp_2d(:,:) = data_slice(:,:,1)
    !call fftwnd_f77_one(fft_2d_id%plan_back, temp_2d(:,:), data_slice(:,:,i3))

    !temp_1d(:) = data_slice(1,1,:)
    !call fftw_f77_one(fft_1d_id%plan_back, temp_1d(:), temp_data(:))

    if(sign == FFT_FORW) then
       do i3=1,n3
          temp_2d = data_slice(:,:,i3)*scale
          call fftwnd_f77_one(fft_2d_id%plan_forw, temp_2d, data_slice(:,:,i3))
       end do
       if(n3 > 1) then
          do i2=1,n2
             do i1=1,n1
                temp_1d = data_slice(i1,i2,:)
                call fftw_f77_one(fft_1d_id%plan_forw, temp_1d, &
                data_slice(i1,i2,:))
             end do
          end do
       end if
    else
       do i3=1,n3
          temp_2d = data_slice(:,:,i3)*scale
          call fftwnd_f77_one(fft_2d_id%plan_back, temp_2d, data_slice(:,:,i3))
       end do
       if(n3 > 1) then
          do i2=1,n2
             do i1=1,n1
                temp_1d = data_slice(i1,i2,:)
                call fftw_f77_one(fft_1d_id%plan_back, temp_1d, &
                data_slice(i1,i2,:))
             end do
          end do
       end if
    end if
    return	
  end subroutine fft_3d_two_pass_data
  !----------------------------------------------------------------------
  subroutine fft_4d_two_pass_data(fft_1d_id,fft_3d_id,sign, &
  n1,n2,n3,n4,data_slice)
    type (fft_type) fft_1d_id,fft_3d_id
    integer :: n1,n2,n3,n4,i1,i2,i3,i4
    integer :: sign
    real    :: scale
    complex,  dimension (:,:,:,:)  :: data_slice
    complex,  dimension (n1,n2,n3) :: temp_3d
    complex,  dimension (n4)       :: temp_1d

    scale=fft_3d_id%scale*fft_1d_id%scale
    if(sign == FFT_FORW) then
       do i4=1,n4
          temp_3d = data_slice(:,:,:,i4)*scale
          call fftwnd_f77_one(fft_3d_id%plan_forw, temp_3d, data_slice(:,:,:,i4))
       end do
       if(n4 > 1) then
          do i3=1,n3
             do i2=1,n2
                do i1=1,n1
                   temp_1d = data_slice(i1,i2,i3,:)
                   call fftw_f77_one(fft_1d_id%plan_forw, temp_1d, &
                   data_slice(i1,i2,i3,:))
                end do
             end do
          end do
       end if
    else   
       do i4=1,n4
          temp_3d = data_slice(:,:,:,i4)*scale
          call fftwnd_f77_one(fft_3d_id%plan_back, temp_3d, data_slice(:,:,:,i4))
       end do
       if(n4 > 1) then
          do i3=1,n3
             do i2=1,n2
                do i1=1,n1
                   temp_1d = data_slice(i1,i2,i3,:)
                   call fftw_f77_one(fft_1d_id%plan_back, temp_1d, &
                   data_slice(i1,i2,i3,:))
                end do
             end do
          end do
       end if
    end if
    return	
  end subroutine fft_4d_two_pass_data
  !----------------------------------------------------------------------
  subroutine clean_fft(fft_id) 
    type (fft_type) fft_id
    integer ierr

    select case (fft_id%ndim)
    case (1)
       call fftw_f77_destroy_plan(fft_id%plan_forw)
       call fftw_f77_destroy_plan(fft_id%plan_back)
    case (2)
       call fftwnd_f77_destroy_plan(fft_id%plan_forw)
       call fftwnd_f77_destroy_plan(fft_id%plan_back)
    case (3)
       call fftwnd_f77_destroy_plan(fft_id%plan_forw)
       call fftwnd_f77_destroy_plan(fft_id%plan_back)
    case default
       write(0,*)'fft_id%ndim=',fft_id%ndim
       call seperr('Not supported fft_id%ndim')
    end select
    return	
  end subroutine clean_fft
!-----------------------------------------------------
  subroutine clean_rfft(fft_id) 
    type (fft_type) fft_id
    integer ierr

    select case (fft_id%ndim)
    case (1)
       call rfftw_f77_destroy_plan(fft_id%plan_forw)
       call rfftw_f77_destroy_plan(fft_id%plan_back)
    case default
       write(0,*)'fft_id%ndim=',fft_id%ndim
       call seperr('Not supported fft_id%ndim')
    end select
    return	
  end subroutine clean_rfft
end module fft_data_mod
