!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! RayleighBernard2D 
! fft.f90 module
! James Sandham 
! 14 Mar 2014
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!


module fft
use def
use mpi
include 'fftw3.f'

integer*8 plan_forward
integer*8 plan_backward

complex, dimension(n) :: scr_in=0.0
complex, dimension(n) :: scr_out=0.0

contains
  subroutine init_fftw_plans()
  call sfftw_plan_dft_1d(plan_forward,n,scr_in,scr_out,FFTW_FORWARD,FFTW_ESTIMATE)
  call sfftw_plan_dft_1d(plan_backward,n,scr_in,scr_out,FFTW_BACKWARD,FFTW_ESTIMATE)
  end subroutine init_fftw_plans

  
  subroutine dest_fftw_plans()
  call sfftw_destroy_plan(plan_forward)
  call sfftw_destroy_plan(plan_backward)
  end subroutine dest_fftw_plans


  subroutine fft1(input_array,output_array)
  complex, dimension(bs), intent(in) :: input_array
  complex, dimension(bs), intent(inout) :: output_array
  integer :: i
  
  do i=1,bs/n
    call sfftw_execute_dft(plan_forward,input_array(1+n*(i-1):n*i),output_array(1+n*(i-1):n*i))
  enddo
  end subroutine fft1

  
  subroutine ifft1(input_array,output_array)
  complex, dimension(bs), intent(in) :: input_array
  complex, dimension(bs), intent(inout) :: output_array
  integer :: i
  
  do i=1,bs/n
    call sfftw_execute_dft(plan_backward,input_array(1+n*(i-1):n*i),output_array(1+n*(i-1):n*i))
  enddo
  output_array = output_array*ni
  end subroutine ifft1

  
  subroutine mpitranspose(input_array)
  complex, dimension(bs), intent(inout) :: input_array

  integer :: i,j,k,ioffa,ioffb,nb,nb2,nbk1,nk1
  complex, dimension(bs) :: scr1, scr2

  nb=n/nproc ; nb2=nb*nb

  do k=1,nb
    nbk1 = nb*(k-1)
    do j=0,nproc-1
      ioffa = nb*j+nproc*nbk1
      ioffb = nb2*j+nbk1
      do i=1,nb
        scr1(i+ioffb) = input_array(i+ioffa)
      enddo
    enddo
  enddo

  call MPI_ALLTOALL(scr1,nb2,MPI_COMPLEX,scr2,nb2,MPI_COMPLEX,MPI_COMM_WORLD,ierr)

  !do k=1,nb
  !  nk1=n*(k-1)
  !  do j=1,n
  !    input_array(j+nk1)=scr2(k+nb*(j-1))
  !  enddo
  !enddo

  do k=1,nb
    nk1 = n*(k-1)
    do j=0,nproc-1
      ioffa = k+nb2*j-nb
      ioffb = nb*j+nk1
      do i=1,nb
        input_array(i+ioffb) = scr2(nb*i+ioffa)
      enddo
    enddo
  enddo  
  end subroutine mpitranspose
end module fft


