module propag_mod

  use fft_data_mod, only: fft_type, init_fft_1d, init_rfft_1d, fft_1d_data, rfft_1d_data, clean_fft, clean_rfft
  implicit none

  private
  public :: prop_init, prop_phsh, clean_prop_phsh
  public :: rft_fwd_axis1, rft_bck_axis1

  real,    parameter :: pi=3.1415926535898
  complex, parameter :: i = cmplx(0.,1.)
  integer            :: sign !sign of phaseshift in exponential
  integer            :: nx, nt, nz, nw, nkx
  real               :: dx, dz, dt, dw, dkx
  real, dimension(:) :: w, kx
  type(fft_type)     :: fft_1d_id, rft_1d_id

  ! Bug danger! The common variables should be modified ONLY by init & clean. Or at worst, by a single subroutine only.

  contains

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  function prop_init(nz, dz, oz, nxm, dxm, oxm, nx, dx, ox, nt, dt, ot)

    integer                :: nz, nxm, nx, nt
    real 		   :: dz, oz, dxm, oxm, dx, ox, dt, ot
    integer                :: nw, nkx
    real 		   :: dw, ow, dkx, okx
    real, allocatable	   :: w(:), kx(:)
    integer                :: iw, ikx

    sign=1 ! if(dowcon_as_shot) sign = +1
    rft_1d_id = init_rfft_1d(nt)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    nw = nt/2
    dw = 1./(nt*dt)
    ow = dw				! no zero frequency component
    allocate(w(nw)); w=0.
    do iw = 1, nw
       w(iw) = ow  + (iw - 1) * dw
    end do
    w = 2*pi * w
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    fft_1d_id = init_fft_1d(nx)
    nkx = nx
    allocate(kx(nkx)); kx=0.

    dkx = 1./(nkx*dx)
    ok = 1./(2*dx)-dkx			!k does not increase uniformly!
    do ik=1, nkx
       kx(ik) = mod(okx+(ik-1)*dkx,1./dx)-okx
    end do

    k = 2*pi * k

  end function post3dvz_init
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine prop_phsh(C, v, D)
! phase-shift modeling

    complex, dimension (:,:)    	       :: C !model in (z,x) domain
    complex, allocatable		       :: U(:,:) !wavefield in (w,kx) domain
    real, allocatable   		       :: D(:,:) !wavefield in (t,x) domain
    real, intent(in)			       :: v(nz,nx) ! 0.5*interval velocity
    real                                       :: v0! reference velocity
    real                                       :: sqa !sqrt(kz)
    integer                                    :: iw,ik,it,ix,iz

    allocate(U(nw,nkx),D(nt,nx))
    U = 0
    call fft_1d_data(fft_1d_id, -1, nkx, C)
!---frequency loop
    do iw = 1, nw
!Operations in the wavenumber domain
       do ik = 1, nkx
!Operations in z
          do iz = nz,-1,1
             write(0,*) (nz-iz+1)*100./nz, '% modeled'
             v0 = sum(v(iz,:))/nkx

!---phase shift term
             sqa = w(iw)**2 / v0**2 - k(ik)**2

!---avoid evan
             if ( sqa > 0 ) then
                U(iw,ik) = U(iw,ik) * exp( i * dz * sqrt(sqa) ) + C(ik,iz)
             else
                U(iw,ik) = 0.
             end if
          end do
       end do 
    end do 
    call fft_1d_data(fft_1d_id, 1, nkx, U)
    D = real(rft_bck_axis1(U))
  end subroutine pprop_phsh

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine clean_prop_phsh()

    call clean_rfft(rft_1d_id)
    call clean_fft (fft_1d_id)
    deallocate(k, w)

  end subroutine clean_prop_phsh
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  function rft_fwd_axis1(U)
    real   , dimension(:,:)        , intent(in) :: U
    complex, dimension(size(U,1)/2,  size(U,2)) :: rft_fwd_axis1
    complex, dimension(size(U,1)/2+1)           :: all_freq
    integer                                     :: ix

    do ix = 1, nx
       call rfft_1d_data(rft_1d_id, -1, nt, U(:,ix), all_freq)
       rft_fwd_axis1(:,ix) = all_freq(2:) ! Throw away the zero frequency
    end do
  end function rft_fwd_axis1
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  function rft_bck_axis1(C)
    complex, dimension(:,:), intent(in)        :: C
    real   , dimension(2*size(C,1), size(C,2)) :: rft_bck_axis1
    complex, dimension(1+size(C,1))            :: all_freq
    integer                                    :: ix

    all_freq (1)=0

    do ix = 1, nx
       all_freq(2:) = C(:,ix)
       call rfft_1d_data(rft_1d_id, 1, nt, rft_bck_axis1(:,ix), all_freq)
    end do
  end function rft_bck_axis1
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!



