!!.. this module contains various kinds of one-way extrapolators
!!--------------------------------------------------------------------
!!.. Author: Yaxun Tang, tang@sep.stanford.edu
module wem_extrpl_mod
!  use  base_slofft_mod
!  use  base_fftwv3_mod
use OW_type
use OW_parms_areal
use OW_fftw_areal
  implicit none

contains
  !!---------------------------------------------------------------------------------------------------------
  !!.. Phase-shift extrapolator
  !!---------------------------------------------------------------------------------------------------------
  integer function ssr_phs(iw, wfld_slice, vmod_slice, vref, opsign) result(stat)
    complex, dimension(:)   :: wfld_slice
    real   , dimension(:)   :: vmod_slice
    real                    :: vref
    integer                 :: opsign
    real                    :: wov


    stat = 1
    wov = ws / vref
    call sfftw_dft_2dx_oper(plan_2dxF, wfld_slice, nx, ny) 
    call phase_shift(wfld_slice, kx2, ky2, vref, ws, dz, nx, ny, opsign)
    call sfftw_dft_2dx_oper(plan_2dxI, wfld_slice, nx, ny)
    stat = 0
  end function

  !!---------------------------------------------------------------------------------------------------------
  !!.. Splt-step Fourier extrapolator
  !!---------------------------------------------------------------------------------------------------------
  integer function ssr_ssf(wfld_slice, vmod_slice, kx2, ky2, vref, ws, dz, nx, ny, opsign, plan_2dxF, plan_2dxI) result(stat)
    complex, dimension(:,:) :: wfld_slice
    real   , dimension(:,:) :: vmod_slice
    real   , dimension(:)   :: kx2, ky2
    real                    :: vref
    real                    :: ws
    real                    :: dz
    integer                 :: nx, ny
    integer                 :: opsign
    integer*8               :: plan_2dxF
    integer*8               :: plan_2dxI

    real                    :: wov

    stat = 1

    wov = ws / vref

    call sfftw_dft_2dx_oper(plan_2dxF, wfld_slice, nx, ny)

    call phase_shift(wfld_slice, kx2, ky2, vref, ws, dz, nx, ny, opsign)

    call sfftw_dft_2dx_oper(plan_2dxI, wfld_slice, nx, ny)

    call ssf_correct(wfld_slice, vmod_slice, vref, ws, dz, nx, ny, opsign)

    stat = 0
  end function
  
  !!-----------------------------------------------------------------------------------------------------------------------
  !!.. Fourier finite-difference extrapolator
  !!-----------------------------------------------------------------------------------------------------------------------
  integer function ssr_ffd(wfld_slice, vmod_slice, kx2, ky2, vref, ws, dx, dy, dz, trick, nx, ny, dip, opsign, rwrk1d, cwrk1d, plan_2dxF, plan_2dxI) result(stat)
    complex, dimension(:,:)          :: wfld_slice
    real   , dimension(:,:)          :: vmod_slice
    real   , dimension(:)            :: kx2, ky2
    real                             :: vref
    real                             :: ws
    real                             :: dx, dy, dz, trick
    integer                          :: nx, ny, dip
    integer                          :: opsign
    real   , dimension(nx+ny)        :: rwrk1d
    complex, dimension(nx*5+ny*5)    :: cwrk1d
    integer*8                        :: plan_2dxF
    integer*8                        :: plan_2dxI

    real                             :: alph(5), beta(5)
    real                             :: wov, vavr
    integer                          :: nstep

    vref = minval(vmod_slice)
    vavr = sum(vmod_slice)/nx*ny

    stat = 1

    call ifd_coeffts(alph, beta, nstep, dip)

    call sfftw_dft_2dx_oper(plan_2dxF, wfld_slice, nx, ny)

    call phase_shift(wfld_slice, kx2, ky2, vref, ws, dz, nx, ny, opsign)

    !call ifd_errfltr(wfld_slice, kx2, ky2, alph, beta, trick, ws, vavr, &
    !                 dx, dy, dz, nx, ny, nstep, opsign)

    call sfftw_dft_2dx_oper(plan_2dxI, wfld_slice, nx, ny)

    call ssf_correct(wfld_slice, vmod_slice, vref, ws, dz, nx, ny, opsign)

    call ffd_correct(wfld_slice, vmod_slice, alph, beta, vref, ws, dx, dy, dz, trick, nx, ny, nstep, dip, opsign, rwrk1d, cwrk1d)

    stat = 0
  end function

!  integer function ssr_efd() result(stat)
!  end function

  !!-----------------------------------------------------------------------------------------------------------
  !!.. Frequency domain implicit finite-difference extrapolator
  !!.. Dip upto 15, 45, 65, 79, 80, 87, 89, 90
  !!-----------------------------------------------------------------------------------------------------------
  integer function ssr_ifd(wfld_slice, vmod_slice, kx2, ky2, vref, ws, dx, dy, dz, trick, nx, ny, dip, opsign, rwrk1d, cwrk1d, plan_2dxF, plan_2dxI) result(stat)
    complex, dimension(:,:)          :: wfld_slice
    real   , dimension(:,:)          :: vmod_slice
    real   , dimension(:)            :: kx2, ky2
    real                             :: vref, ws
    real                             :: dx, dy, dz, trick
    integer                          :: nx, ny, dip
    integer                          :: opsign
    real   , dimension(nx+ny)        :: rwrk1d
    complex, dimension(nx*5+ny*5)    :: cwrk1d

    integer*8                        :: plan_2dxF
    integer*8                        :: plan_2dxI

    integer                          :: ix, iy, istep, nstep
    integer                          :: lrx1, lry1
    integer                          :: lcx1, lcx2, lcx3, lcx4, lcx5
    integer                          :: lcy1, lcy2, lcy3, lcy4, lcy5
    real                             :: alph(5), beta(5)
    

!integer, save :: counter=0

    !!.. asign work space
    lcx1 = 1            !!.. wfdx
    lcx2 = lcx1 + nx    !!.. ahx
    lcx3 = lcx2 + nx    !!.. bhx
    lcx4 = lcx3 + nx    !!.. chx
    lcx5 = lcx4 + nx    !!.. rhx
    lcy1 = lcx5 + nx    !!.. wfdy
    lcy2 = lcy1 + ny    !!.. ahy
    lcy3 = lcy2 + ny    !!.. bhy
    lcy4 = lcy3 + ny    !!.. chy
    lcy5 = lcy4 + ny    !!.. rhy
    
    lrx1 = 1            !!.. velx
    lry1 = lrx1 + nx    !!.. vely

    call ifd_coeffts(alph, beta, nstep, dip)

    if (nx > 1) then
       do iy=1, ny
          cwrk1d(lcx1:lcx2-1) = wfld_slice(:,iy)
          rwrk1d(lrx1:lry1-1) = vmod_slice(:,iy)
          do istep=1, nstep
             call ifd_sysform(cwrk1d(lcx2), cwrk1d(lcx3), cwrk1d(lcx4), rwrk1d(lrx1), ws, dx, dz, alph(istep), beta(istep), trick, nx, opsign)
             call ifd_dfrterm(cwrk1d(lcx1), cwrk1d(lcx2), cwrk1d(lcx3), cwrk1d(lcx4), cwrk1d(lcx5), nx)
          end do
          wfld_slice(:,iy) = cwrk1d(lcx1:lcx2-1)
       end do
    end if

    if (ny > 1) then
       do ix=1, nx
          cwrk1d(lcy1:lcy2-1) = wfld_slice(ix,:)
          rwrk1d(lry1:nx+ny ) = vmod_slice(ix,:)
          do istep=1, nstep
             call ifd_sysform(cwrk1d(lcy2), cwrk1d(lcy3), cwrk1d(lcy4), rwrk1d(lry1), ws, dy, dz, alph(istep), beta(istep), trick, ny, opsign)
             call ifd_dfrterm(cwrk1d(lcy1), cwrk1d(lcy2), cwrk1d(lcy3), cwrk1d(lcy4), cwrk1d(lcy5), ny)
          end do
          wfld_slice(ix,:) = cwrk1d(lcy1:lcy2-1)
       end do
    end if

    call ifd_sftterm(wfld_slice, vmod_slice, ws, dz, nx, ny, opsign)

!counter = counter + 1
!    if (counter/8*8==counter) then
    !!.. Li error compensation filter
!    call sfftw_dft_2dx_oper(plan_2dxF, wfld_slice, nx, ny)

!    call ifd_errfltr(wfld_slice, kx2, ky2, alph, beta, trick, ws, vref, &
!                     dx, dy, dz, nx, ny, nstep, opsign)

!    call sfftw_dft_2dx_oper(plan_2dxI, wfld_slice, nx, ny)
!    end if
    stat = 0
  end function

  !!-------------------------------------------------------------------------------------------
  !!.. solving for diffraction term
  !!------------------------------------------------------------------------------------------- 
  subroutine ifd_dfrterm(wfdh, ah, bh, ch, rh, nh)
    complex, dimension(nh) :: wfdh
    complex, dimension(nh) :: ah, bh, ch, rh
    integer                :: nh

    integer                :: ih

    rh = 0.

    do ih=2, nh-1
       rh(ih) = conjg(ah(ih))*wfdh(ih-1) &
              + conjg(bh(ih))*wfdh(ih  ) &
              + conjg(ch(ih))*wfdh(ih+1)
    end do

    !bh(1)  = -ch( 1)
    !bh(nh) = -ah(nh)

    rh(1)  = conjg( bh(1))*wfdh(   1) &
           + conjg( ch(1))*wfdh(   2)

    rh(nh) = conjg(ah(nh))*wfdh(nh-1) &
           + conjg(bh(nh))*wfdh(nh  )

    call ctridag(ah, bh, ch, rh, wfdh, nh)

    return
  end subroutine
  !!--------------------------------------------------------------------------
  !!.. solving for shifting term
  !!--------------------------------------------------------------------------
  subroutine ifd_sftterm(wfld_slice, vmod_slice, ws, dz, nx, ny, opsign)
    complex, dimension(:,:) :: wfld_slice
    real   , dimension(:,:) :: vmod_slice
    real                    :: ws, dz
    integer                 :: nx, ny
    integer                 :: opsign

    integer                 :: ix, iy
    real                    :: wov, shift

    do iy=1, ny
    do ix=1, nx
       wov = ws / vmod_slice(ix,iy) 
       shift = opsign * wov * dz
       wfld_slice(ix,iy) = wfld_slice(ix,iy) * cmplx(cos(shift), sin(shift))
    end do 
    end do
    return
  end subroutine


  !!---------------------------------------------------------------------------
  !!.. finite-difference error compensation filter
  !!.. Li, Z., Compensating finite-difference errors in 3-D migration and 
  !!.. modeling, 1991, VOL. 56, Geophysics. 
  !!---------------------------------------------------------------------------
  subroutine ifd_errfltr(wfld_slice, kx2, ky2, alph, beta, trick, ws, vref, &
                         dx, dy, dz, nx, ny, nstep, opsign)
    complex, dimension(:,:) :: wfld_slice
    real   , dimension(:)   :: kx2, ky2
    real   , dimension(5)   :: alph, beta
    real                    :: trick
    real                    :: ws, vref
    real                    :: dx, dy, dz 
    integer                 :: nx, ny, nstep, opsign

    integer                 :: ix, iy, istep
    real                    :: Sx, Sy, wov, spltx, splty
    real                    :: kz2, kz, kn, phase

    wov = ws / vref 

    do iy=1, ny
    do ix=1, nx
       Sx    = - kx2(ix)/wov/wov
       Sy    = - ky2(iy)/wov/wov
       kz2   = 1. + Sx + Sy

       spltx = 0.
       splty = 0. 

       do istep=1, nstep
          spltx = spltx + alph(istep)*Sx/(1+trick*Sx)/(1.+beta(istep)*Sx/(1+trick*Sx))
          splty = splty + alph(istep)*Sy/(1+trick*Sy)/(1.+beta(istep)*Sy/(1+trick*Sy))
       end do

       kn = 1. + spltx + splty

       if (kz2>=0) then
          kz    = sqrt(kz2)
          phase = opsign*wov*(kz - kn)*dz
          wfld_slice(ix,iy) = wfld_slice(ix,iy)*cmplx(cos(phase), sin(phase))
       else
          kz    = sqrt(-kz2)
          phase = - wov*kz*dz
          wfld_slice(ix,iy) = wfld_slice(ix,iy)*exp(phase)
       end if

    end do
    end do

    return
  end subroutine

  !!---------------------------------------------------------------------------------------
  !!.. Forming the tridiagonal matrix for implicit finite-difference extrapolation
  !!---------------------------------------------------------------------------------------
  subroutine ifd_sysform(ah, bh, ch, velh, ws, dh, dz, alph, beta, trick, nh, opsign)
    complex, dimension(nh) :: ah, bh, ch
    real   , dimension(nh) :: velh
    real                   :: ws, dh, dz
    real                   :: alph, beta, trick
    integer                :: nh, opsign

    integer                :: ih
    real                   :: wov
    complex                :: aah, bbh

    ah = 0; bh = 0; ch = 0

    do ih=1, nh
       wov = ws/velh(ih)
       aah = opsign * alph*cmplx(0.,1.)*dz/2./wov/dh/dh
       bbh = beta/dh/dh/wov/wov + trick
    
       ah(ih) = bbh - aah            !!.. sub-diagonal   elements: 2~N 
       ch(ih) = bbh - aah            !!.. super-diagonal elements: 1~N-1  
       bh(ih) = 1 - 2.*(bbh - aah)   !!.. diagonal elements
    end do 

    return 
  end subroutine

  !!----------------------------------------------------------------------------------------------
  !!.. forming the tridiagonal matrix for FFD extrapolation
  !!----------------------------------------------------------------------------------------------
  subroutine ffd_sysform(ah, bh, ch, velh, vref, ws, dh, dz, alph, beta, trick, nh, opsign)
    complex, dimension(nh) :: ah, bh, ch
    real   , dimension(nh) :: velh
    real                   :: vref, ws, dh, dz
    real                   :: alph, beta, trick
    integer                :: nh, opsign

    integer                :: ih
    real                   :: wov, rov
    complex                :: aah, bbh

    ah = 0; bh = 0; ch = 0

    do ih=1, nh
       wov = ws/velh(ih)
       rov = vref/velh(ih)  
       aah = opsign * alph*cmplx(0.,1.)*dz/2./wov/dh/dh*(1-rov)
       bbh = beta/dh/dh/wov/wov*(rov*rov+rov+1.) + trick

       ah(ih) = bbh - aah
       ch(ih) = bbh - aah
       bh(ih) = 1 - 2.*(bbh - aah)
    end do

    return
  end subroutine

  !!-----------------------------------------------------------------------
  !!.. Phase-shift kernel
  !!-----------------------------------------------------------------------
  subroutine phase_shift(wfld_slice, kx2, ky2, vref, ws, dz, nx, ny, opsign)
    complex, dimension(:,:) :: wfld_slice
    real   , dimension(:)   :: kx2
    real   , dimension(:)   :: ky2
    real                    :: vref, ws
    real                    :: dz
    integer                 :: nx, ny
    integer                 :: opsign

    integer                 :: ix, iy
    real                    :: kz, kz2
    real                    :: phase, wov

    wov = ws / vref

    do iy=1, ny
    do ix=1, nx
       kz2 = 1. - kx2(ix)/wov/wov - ky2(iy)/wov/wov
       if (kz2 >= 0.) then
          kz = sqrt(kz2)
          phase = opsign*wov*kz*dz 
          wfld_slice(ix,iy) = wfld_slice(ix,iy)*cmplx(cos(phase),sin(phase))
       else
          kz = sqrt(-kz2)
          phase = - wov*kz*dz
          wfld_slice(ix,iy) = wfld_slice(ix,iy)*exp(phase)
       end if
    end do
    end do
    return
  end subroutine


  !!-------------------------------------------------------------------------
  !!.. Split-step correction
  !!-------------------------------------------------------------------------
  subroutine ssf_correct(wfld_slice, vmod_slice, vref, ws, dz, nx, ny, opsign)
    complex, dimension(:,:) :: wfld_slice
    real   , dimension(:,:) :: vmod_slice
    real                    :: vref
    real                    :: ws
    real                    :: dz
    integer                 :: nx, ny
    integer                 :: opsign
 
    integer                 :: ix, iy
    real                    :: split, wov

    wov = ws / vref

    do iy=1, ny
    do ix=1, nx
       split = opsign*wov*dz*(vref/vmod_slice(ix,iy)-1.)
       wfld_slice(ix,iy) = wfld_slice(ix,iy)*cmplx(cos(split),sin(split))
    end do 
    end do
    return
  end subroutine

  !!--------------------------------------------------------------------------
  !!.. Fourier finite-difference correction
  !!--------------------------------------------------------------------------
  subroutine ffd_correct(wfld_slice, vmod_slice, alph, beta, vref, ws, dx, dy, dz, trick, nx, ny, nstep, dip, opsign, rwrk1d, cwrk1d)
    complex, dimension(:,:)          :: wfld_slice
    real   , dimension(:,:)          :: vmod_slice
    real                             :: vref, ws
    real                             :: dx, dy, dz, trick
    integer                          :: nx, ny, dip, opsign
    real   , dimension(nx+ny)        :: rwrk1d
    complex, dimension(nx*5+ny*5)    :: cwrk1d

    real                             :: alph(5), beta(5)
    integer                          :: ix, iy, istep, nstep 
    integer                          :: lrx1, lry1
    integer                          :: lcx1, lcx2, lcx3, lcx4, lcx5
    integer                          :: lcy1, lcy2, lcy3, lcy4, lcy5

    !!.. asign work space
    lcx1 = 1           !!.. wfdx
    lcx2 = lcx1 + nx    !!.. ahx
    lcx3 = lcx2 + nx    !!.. bhx
    lcx4 = lcx3 + nx    !!.. chx
    lcx5 = lcx4 + nx    !!.. rhx
    lcy1 = lcx5 + nx    !!.. wfdy
    lcy2 = lcy1 + ny    !!.. ahy
    lcy3 = lcy2 + ny    !!.. bhy
    lcy4 = lcy3 + ny    !!.. chy
    lcy5 = lcy4 + ny    !!.. rhy

    lrx1 = 1
    lry1 = lrx1 + nx

    !call ifd_coeffts(alph, beta, nstep, dip)

    if (nx > 1) then
       do iy=1, ny
          cwrk1d(lcx1:lcx2-1) = wfld_slice(:,iy)
          rwrk1d(lrx1:lry1-1) = vmod_slice(:,iy)
          do istep=1, nstep
             call ffd_sysform(cwrk1d(lcx2), cwrk1d(lcx3), cwrk1d(lcx4), rwrk1d(lrx1), vref, ws, dx, dz, alph(istep), beta(istep), trick, nx, opsign)
             call ifd_dfrterm(cwrk1d(lcx1), cwrk1d(lcx2), cwrk1d(lcx3), cwrk1d(lcx4), cwrk1d(lcx5), nx)
          end do
          wfld_slice(:,iy) = cwrk1d(lcx1:lcx2-1)
       end do
    end if

    if (ny > 1) then
       do ix=1, nx
          cwrk1d(lcy1:lcy2-1) = wfld_slice(ix,:)
          rwrk1d(lry1:nx+ny ) = vmod_slice(ix,:)
          do istep=1, nstep
             call ffd_sysform(cwrk1d(lcy2), cwrk1d(lcy3), cwrk1d(lcy4), rwrk1d(lry1), vref, ws, dy, dz, alph(istep), beta(istep), trick, ny, opsign)
             call ifd_dfrterm(cwrk1d(lcy1), cwrk1d(lcy2), cwrk1d(lcy3), cwrk1d(lcy4), cwrk1d(lcy5), ny)
          end do
          wfld_slice(ix,:) = cwrk1d(lcy1:lcy2-1)
       end do
    end if

    return
  end subroutine
  !!---------------------------------------------------------------------------
  !!.. complex tridiagonal solver
  !!---------------------------------------------------------------------------
  subroutine ctridag(a,b,c,r,u,n)
    integer            :: n
    complex            :: a(n), b(n), c(n), r(n), u(n)
    integer            :: j
    integer, parameter :: NMAX = 5000
    complex            :: bet, gam(NMAX)
    if (b(1) == 0.) pause "ctirdag: rewrite equations"
    bet  = b(1)
    u(1) = r(1)/bet
    do j=2, n
       gam(j) = c(j-1)/bet
       bet    = b(j)-a(j)*gam(j)
       if (bet == 0.) pause "ctridag failed"
       u(j) = (r(j)-a(j)*u(j-1))/bet
    end do 
    do j=n-1,1,-1
       u(j) = u(j) - gam(j+1)*u(j+1)
    end do
    return
  end subroutine


  !!-----------------------------------------------------------------------------------------
  !!.. set coefficients for implicit finite-difference extrapolator
  !!.. based on Zaitian Ma higher order splitting method
  !!.. Ma, Z., 1981, Finite-difference migration with higher-order approximation:
  !!.. Presented at the 1981 joint meeting of the China Geophysical Society and
  !!.. Soc. of Expl. Geophys., Beijing, China
  !!-----------------------------------------------------------------------------------------
  subroutine ifd_coeffts(alph, beta, nstep, dip)
    real    :: alph(5), beta(5)
    integer :: nstep
    integer :: dip
    select case (dip)
      case (15)
           alph(1) = 0.500000000; beta(1) = 0.000000000
           nstep = 1

      case (45)
           alph(1) = 0.500000000; beta(1) = 0.250000000
           nstep = 1

      case (65)
           alph(1) = 0.478242060; beta(1) = 0.376369527
           nstep = 1

      case (79)
           alph(1) = 0.457500000; beta(1) = 0.457500000
           nstep = 1

      case (80)
           alph(1) = 0.457289566; beta(1) = 0.222691983
           alph(2) = 0.040315157; beta(2) = 0.873981642
           nstep = 2

      case (87)
           alph(1) = 0.414236605; beta(1) = 0.150843924
           alph(2) = 0.081312882; beta(2) = 0.744418059
           alph(3) = 0.004210420; beta(3) = 0.972926132
           nstep = 3

      case (89)
           alph(1) = 0.367013245; beta(1) = 0.105756624
           alph(2) = 0.117592008; beta(2) = 0.614520676
           alph(3) = 0.014853510; beta(3) = 0.919432661
           alph(4) = 0.000523275; beta(4) = 0.994065088
           nstep = 4

      case (90)
           alph(1) = 0.318013812; beta(1) = 0.073588213
           alph(2) = 0.143798076; beta(2) = 0.483340757
           alph(3) = 0.033860918; beta(3) = 0.824918565
           alph(4) = 0.004172967; beta(4) = 0.964827992
           alph(5) = 0.000153427; beta(5) = 0.997370236
           nstep = 5

    end select
    return
  end subroutine



end module
