module weiop_wem
  !! Wave-Equation Migration
  !! full 3-D prestack

  use wei_util
  use wei_pro
  use wei_par
  use wei_kmap
  use wei_comp
  use wei_attr
  !---------------
  implicit none

  real,                           private :: dkxko
  integer,                        private :: nkxko
  real,                           private ::  kxko

  complex, dimension(:,:), pointer, private ::  kzko, tkzko
contains
  !----------------------------------------------------------------
  subroutine wemdsr_init(max,n)
    real,    intent(in) :: max
    integer, intent(in) :: n
    integer             :: er
    integer             :: i,m
    real                :: kxko
    integer             :: iomp
    character(len=128), parameter :: unit='wemdsr_init'

    call weimem_add(8,n*wpar%nomp,unit)
    allocate( kzko(   n,wpar%nomp),stat=er)
    if(er/=0) call seperr('cant allocate  kzko')

    call weimem_add(8,n*wpar%nomp,unit)
    allocate(tkzko(   n,wpar%nomp),stat=er)
    if(er/=0) call seperr('cant allocate tkzko')

    dkxko = max/(n-1)
    nkxko = n
    m=n/max

    do iomp=1,wpar%nomp
       !!                    ____________
       !!   kz         1    /     | kx |2          kx
       !! i -- dz = -i -   /  1 - | -- |  dz   0 < -- < 1 
       !!   ko         2 \/       | ko |           ko
       !!
       do i=1,m
          kxko    = 0. +     1.0 * (i  -1)/(  m-1) !! kx/ko=0...1
          kzko(i,iomp) = - (0,+1) * sqrt(1-kxko**2)
       end do

       !!                    ____________
       !!   kz         1    / | kx |2              kx
       !! i -- dz = -  -   /  | -- | - 1  dz   1 < -- < max 
       !!   ko         2 \/   | ko |               ko
       !!
       do i=m+1,n
          kxko    = 1. + (max-1) * (i-m-1)/(n-m-1) !! kx/ko=1...max
          kzko(i,iomp) = - (+1,0) * sqrt(kxko**2-1)
       end do

       kzko(:,iomp) = .5 * kzko(:,iomp)
    end do

  end subroutine wemdsr_init
  !----------------------------------------------------------------
  integer function weiop_wem3_init() result(st)
    integer :: er
    integer :: mpierr
    character(len=128), parameter :: unit='weiop_wem3_init'
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    !! init DSR
    if(wpar%impi==0) call from_param("nkxko",nkxko,2001)
    if(wpar%nmpi >0) call MPI_BCAST(nkxko, 1,MPI_INTEGER,0,MPI_COMM_WORLD,mpierr)
    if(wpar%impi==0) call from_param( "kxko", kxko,2.)
    if(wpar%nmpi >0) call MPI_BCAST( kxko, 1,MPI_REAL   ,0,MPI_COMM_WORLD,mpierr)
    call wemdsr_init(kxko,nkxko)


#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    st=weiOK
  end function weiop_wem3_init
  !----------------------------------------------------------------
  integer function weiop_wem3(iws,izs,ifk,ith,wfld) result(st)
    integer, intent(in) ::iws,izs,ifk,ith
    complex, dimension(:,:,:,:,:), pointer :: wfld
    real    :: w, ko
    integer :: imx,imy,ihx,ihy 
    integer :: ii
    complex :: ikz
    real    :: kk
    real,pointer :: kx(:,:),ky(:,:)
    character(len=128), parameter :: unit='weiop_wem3'
#ifdef DEBUG
    if(wpro%dbg>9) call in(unit)
#endif

    w  = wpro%bw%o + (iws-1)*wpro%bw%d
    ko = 2*w* srp(ith)%rs(izs,ifk)


    if(wpro%opsign==weiB2T) then
       tkzko(:,ith)=conjg(kzko(:,ith))
    else
       tkzko(:,ith)=      kzko(:,ith)
    end if

    do ihy=1,wpro%ahy%n
       do ihx=1,wpro%ahx%n
          do imy=1,wpro%amy%n
             do imx=1,wpro%amx%n

                kk = sqrt(srp(ith)%kx(imx,ihx)**2 + srp(ith)%ky(imy,ihy)**2)
                ii = max(1, min(int(1 + kk/ko / dkxko) , nkxko))
                ikz= ko*tkzko(ii,ith)

                wfld(imx,imy,ihx,ihy,iws) = &
                wfld(imx,imy,ihx,ihy,iws) * &
                cexp( ikz * wpro%az%d)

             end do
          end do
       end do
    end do

#ifdef DEBUG
    if(wpro%dbg>9) call out(unit)
#endif
    st=weiOK
  end function weiop_wem3
  !----------------------------------------------------------------
end module weiop_wem





