!! wavefield extrapolation via Split-Step Fourier for one frequency and one depth step
!! should initialize in advance for each frequency and depth step 
module wei_ssf_mod
  use  wei_myfft_mod
  use  wei_taper_mod
  implicit none

  real, dimension(:), pointer, private :: kx2
  real, dimension(:), pointer, private :: ky2 
  integer                    , private :: dz 
  real                       , private :: kz2_max
contains

!=========================================================================================
! 
  integer function wei_ssf_init(kx2_in, ky2_in, dz_in, maxang_in) result(stat)
    real, dimension(:), pointer :: kx2_in
    real, dimension(:), pointer :: ky2_in
    real                        :: dz_in
    real                        :: maxang_in
    
    stat = 1

    kx2     => kx2_in
    ky2     => ky2_in
    dz      =  dz_in
    kz2_max =  cos(maxang_in*4.*atan(1.)/180.)**2

    stat = 0

  end function

!=========================================================================================
  integer function wei_ssf_clos() result(stat)
    if(associated(kx2)) nullify(kx2)
    if(associated(ky2)) nullify(ky2)
    dz      = 0.
    kz2_max = 0.
    stat = 0
  end function
!=========================================================================================
! cmod(nkx, nky), cdat(nkx, nky) 
  integer function wei_ssf_oper(adj, opsign, ws, vref, vmod_slice, wfld_slice, vzref) result(stat)
    optional                         :: vzref
    logical                          :: adj
    integer                          :: opsign
    real                             :: ws
    real                             :: vref
    real                             :: vzref
    real   , dimension(:,:), pointer :: vmod_slice
    complex, dimension(:,:), pointer :: wfld_slice


    stat = 1

    if (adj) then
       call FFT2DX(wfld_slice, 1)
       !call ssr_phase_shift(opsign, ws, vref, wfld_slice, vzref)
       call ssr_phase_shift(opsign, ws, vref, wfld_slice)
       call FFT2DX(wfld_slice, -1)
       call ssr_split_step( opsign, ws, vref, vmod_slice, wfld_slice)
    else
       call ssr_split_step( opsign, ws, vref, vmod_slice, wfld_slice)
       call FFT2DX(wfld_slice, 1)
       !call ssr_phase_shift(opsign, ws, vref, wfld_slice, vzref)
       call ssr_phase_shift(opsign, ws, vref, wfld_slice)
       call FFT2DX(wfld_slice, -1)
    end if
    stat = 0

  end function

!=====================================================================
!
  subroutine ssr_phase_shift(opsign, ws, vref, wfld_slice, vzref)
    optional                         :: vzref
    integer                          :: opsign
    real                             :: ws
    real                             :: vref
    real                             :: vzref
    complex, dimension(:,:), pointer :: wfld_slice

    real                             :: wov, wov2, kz, kz2, phase, ampli
    integer                          :: ikx, iky


    wov  = ws / vref
    wov2 = wov * wov
    do iky=1, size(ky2)
       do ikx=1, size(kx2)
          kz2 = (1.0 - kx2(ikx)/wov2 - ky2(iky)/wov2)
 
          if (kz2 >= 0. .and. kz2 < kz2_max) then
             wfld_slice(ikx, iky) = 0.

          else if(kz2 >= kz2_max) then
             kz = sqrt(kz2)
             phase = opsign * wov * kz * dz
             wfld_slice(ikx, iky) = wfld_slice(ikx, iky) * cmplx(cos(phase), sin(phase))
          else
             kz = sqrt(-kz2)
             phase = - wov * kz * dz
             wfld_slice(ikx, iky) = wfld_slice(ikx, iky) * exp(phase)
          end if

          if (present(vzref)) then
             ampli = (1. + (kx2(ikx) + ky2(iky))/(wov2 - kx2(ikx) - ky2(iky))) * vzref / (2. * vref)
             wfld_slice(ikx, iky) = wfld_slice(ikx, iky) * exp(-ampli)
          end if

       end do
    end do
  end subroutine

!============================================================================
!
  subroutine ssr_split_step(opsign, ws, vref, vmod_slice, wfld_slice)
    integer                          :: opsign
    real                             :: ws
    real                             :: vref
    real   , dimension(:,:), pointer :: vmod_slice
    complex, dimension(:,:), pointer :: wfld_slice

    integer                          :: ix, iy
    real                             :: split


    do iy=1, size(ky2)
       do ix=1, size(kx2)
          split = opsign * ws * dz * (1./vmod_slice(ix,iy) - 1./vref) 
          wfld_slice(ix, iy) = wfld_slice(ix, iy) * cmplx(cos(split), sin(split))
       end do
    end do
  end subroutine
!=============================================================================
  subroutine ssr_boundary_condition(ws, vref, wfld_slice)
    real                             :: ws
    real                             :: vref
    complex, dimension(:,:), pointer :: wfld_slice  
    complex                          :: shaping
    integer                          :: ikx, iky
    real                             :: wov, wov2, kz, kz2, lamda
    wov  = ws / vref
    wov2 = wov * wov
    call FFT2DX(wfld_slice,  1)
    do iky=1, size(ky2)
       do ikx=1, size(kx2)
          kz2 = (1.0 - kx2(ikx)/wov2 - ky2(iky)/wov2)
          if (kz2 >= 0. .and. kz2 < kz2_max) then
             wfld_slice(ikx, iky) = 0.

          else if(kz2 >= kz2_max) then
             kz = sqrt(kz2)
             lamda = wov * kz
             shaping = cmplx(0., 0.5 / lamda)
             wfld_slice(ikx, iky) = wfld_slice(ikx, iky) * shaping

          else
             kz = sqrt(-kz2)
             lamda = wov * kz 
             shaping = cmplx(0., 0.5 / lamda)
             wfld_slice(ikx, iky) = wfld_slice(ikx, iky) * shaping

          end if
       end do
    end do
    call FFT2DX(wfld_slice, -1)

  end subroutine

end module














