module wem_mod
  !! Wave-Equation propagation
  !! full 3-D prestack SSR

  use util_mod
  use down_mod
  use external_mod
  use var_mod
  use run_mod
  use S_mod
  use kmap_mod
  use sep
  !---------------
  implicit none

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

  complex, dimension(:,:), allocatable, private :: sc
  complex, dimension(:)  , allocatable, private :: kzko
  real   , dimension(:,:), allocatable, private :: bin,bout
contains
  !----------------------------------------------------------------
!  subroutine wemssr_init(max,n)
  subroutine wemssr_init()
    real    :: max
    integer :: n
    integer             :: er
    integer             :: i,m,r
    real                :: kxko
    call from_param("nkxko",n,2001)
    call from_param( "kxko", max,2.)
!    call from_param('maxang',maxang,90.)



    dkxko = max/(n-1)
    nkxko = n
    m=n/2.
write(0,*) 'max,m ,n,d',max,m ,n,dkxko
write(0,*) 'old wem kx ky'
    !!                    ____________
    !!   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) = - (0.,+1.) * sqrt(1-kxko**2)
    end do
write(0,*)' second half'
    !!                    ____________
    !!   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) = - (+1.,0.) * sqrt(kxko**2-1)
    end do

  end subroutine wemssr_init
  !----------------------------------------------------------------
  integer function wem_init() result(st)
    integer :: er
    integer ::i,imax,j
    real :: hxmax,hymax,a
#ifdef DEBUG
    character(len=128), parameter :: unit='wem_init'
    if(run%debug) call in(unit)
#endif

    call from_param("nkxko",nkxko,2001)
    call from_param( "kxko", kxko,2.)

    allocate( kzko(nkxko),stat=er)
    if(er/=0) call seperr('cant allocate  kzko: wem.F90')
    allocate(bin (down%amx%n,run%o_nth),stat=er)
    if(er/=0) call seperr('cant allocate  bin: wem.F90')
    allocate(bout(down%amx%n,run%o_nth),stat=er)
    if(er/=0) call seperr('cant allocate  bout:wem.F90')
    allocate(sc  (down%amx%n,run%o_nth),stat=er)
    if(er/=0) call seperr('cant allocate  sc:wem.F90')
#ifdef MEMCHECK
    a=8*(nkxko+down%amx%n*run%o_nth)
    a=a+4*(down%amx%n*run%o_nth+down%amx%n*run%o_nth)
    a=a/(1024*1024)
    down%mem=down%mem+a
    write(0,*) 'wem.F90: allocate',a,'Mbytes.  Total=',down%mem
#endif

!    call wemssr_init(kxko,nkxko)
write(0,*)' weminit befre'
    call wemssr_init()
write(0,*)' wem bater'
    if(down%opsign==goB2T) then
       do i=1,size(kzko,1)
          kzko(i)=conjg(kzko(i))
       end do
    end if
write(0,*)' kmap',kmap%kx
write(0,*) 'y',kmap%ky

#ifdef DEBUG
    if(run%debug) call out(unit)
#endif
    st=OK
  end function wem_init
!----------------------------------------------------------------
  integer function wem(iws,izs,ivel,ith,wfld) result(st)
    integer, intent(in) ::iws,izs,ivel,ith
    complex, dimension(:,:,:), pointer :: wfld
    real    :: w, ko, k
    integer :: ix,iy,is,ir,i,iz
    complex :: ikz
#ifdef DEBUG
    character(len=128), parameter :: unit='wem'
    if(run%depar) call in(unit)
#endif

    w  = down%bw%o + (iws-1)*down%bw%d
    ko = slo%sz(izs,ivel) * w
!write(0,*)' wem',down%az%d,w,ko
    ky_loop:do iy=1,down%amy%n
       kx_loop:do ix=1,down%amx%n
          k = sqrt(kmap%kx(ix)**2 + kmap%ky(iy)**2)
          i = max(1, min(int(1 + k/ko / dkxko) , nkxko))
!          ikz = ko*kzko(i)*var%dzb(izs)
          ikz = ko*kzko(i)*down%az%d
          bin(ix,ith)=real(ikz)
          sc(ix,ith)=cmplx(cos(aimag(ikz)),sin(aimag(ikz)))
       end do kx_loop
       do ix=1,size(bout,1)
          bout(ix,ith)=exp(bin(ix,ith))
       end do
       do ix=1,size(bout,1)
          wfld(ix,iy,iws)=wfld(ix,iy,iws)*bout(ix,ith)*sc(ix,ith)
       end do
    end do ky_loop

#ifdef DEBUG
    if(run%depar) call out(unit)
#endif
    st=OK
  end function wem
end module wem_mod