module RWE2D_SI_wemig
  use sep
  use util_
  use fft

  implicit none
  integer, private :: nz,nx,nw,nr
  real,    private :: dz,dx,dw
  real,    private :: oz,ox,ow

  integer          :: ntap
  integer, private :: nloop,nref

  real   , private :: forward,eps,pi
  real   , private :: okx,dkx
  logical, private :: kin,verbose

  real,    allocatable,dimension(:),private :: msk,mtt,tap
  complex, allocatable,dimension(:),private :: wfl,wtt,ikz

contains
  !----------------------------------------------------------------  
  subroutine wemig_init(ax,az,aw,ar,forward_in,kin_in,nref_in,verbose_in)
    type(myaxis)     :: ax,az,aw,ar
    integer          ::                               nref_in
    real             ::             forward_in
    logical          ::                        kin_in,        verbose_in
    real             :: pi
    pi=acos(-1.)
    forward=forward_in
    nz=az%n;    nx=ax%n;   nw=aw%n
    dz=az%d;    dx=ax%d;   dw=aw%d * 2.*pi
    oz=az%o;    ox=ax%o;   ow=aw%o * 2.*pi

    nref=nref_in

    call from_param("ntap",ntap,5)
    call from_param("nloop",nloop,3)
    okx=  -pi/    dx
    if (nx .eq. 1) okx=0.

    dkx=2.*pi/(nx*dx)

    kin=kin_in
    verbose=verbose_in;    
    eps=0.0000001

    allocate(ikz(nx))
    allocate(wfl(nx))
    allocate(wtt(nx))
    allocate(msk(nx)); msk=1.
    allocate(mtt(nx)); mtt=1.
    allocate(tap(nx)); tap=1.
    call rwetaper_init()

  end subroutine wemig_init

  !----------------------------------------------------------------

  !! WE migration driver
  subroutine wemig(rwf,ishot,fields,refs,mask)
    integer          :: iw,iz,ix,iy,ir,ishot,mask(:,:)
    real             :: w,refs(:,:,:),fields(:,:,:)
    complex, pointer :: rwf (:,:),rax(:)

    do iw=1,nw  !! Frequency Loop           
       w=ow+(iw-1)*dw
       rax => rwf(:,iw) 
       call sep_write(rax,"SP",nx,8)
       do iz=2,nz  !! Depths
     
          wtt=rax;   rax=0. 
          do ir=1,nref  
             if (refs(ir,3,iz) .le. 10e-10) cycle 
             wfl=wtt
             call wemig_phs(wfl,w,iz,ir,               refs(:,:,iz),1.)
             call wemig_ssf(wfl,w,iz,ir,fields(:,:,iz),refs(:,:,iz),1.)       
             call wemig_ref(wfl,iz,ir,mask(:,iz),rax)
          end do
          call rwetaper(rax)

       !if (mod (iz,10) .eq. 10) 
       call sep_write(rax,"SP",nx,8)

       end do !! . . z-loop
!       write(0,*) maxval(abs(rax)),minval(abs(rax))
       if( mod(iw,10) .eq. 0) write(0,*) 'FINISHED FREQ ',iw,' OF ',nw 

    end do  !! . . w-loop
 

  end subroutine wemig

  !----------------------------------------------------------------

  subroutine wemig_ref(wfl,iz,ir,allmask,dax)
    complex, pointer :: dax(:)
    complex          :: wfl(:)
    integer          :: iz,ir,ix,ii,allmask(:)

    where( allmask .eq. ir) 
       mtt=1.
    elsewhere
       mtt=0.
    end where

    do ii=1,nloop !! Smoothing Iterations
       msk(1) = mtt(1)
       msk(nx)=mtt(nx-1)
       do ix=2,nx-1
          msk(ix)=( mtt(ix-1)+mtt(ix+1) )/2.
       end do
       mtt=msk
    end do

    dax = dax + msk*wfl
    !  write(0,*) minval(abs(wfl)),maxval(abs(wfl))
  end subroutine wemig_ref

  !----------------------------------------------------------------  

  !! Fourier-domain phase shift
  subroutine wemig_phs(dax,w,iz,ir,allrefs,dir)
    integer :: ix,iy,iz,ir
    real    :: w,kx,ky,dir,allrefs(:,:)
    complex :: kz,carg,dax(:)

    call fth(.false.,.false.,dax)

    !! Loop over all kx
    do ix=1,nx
       kx=okx+(ix-1)*dkx

       !! Evaluate the complex exponential argument
       carg = &
         allrefs(ir,3 )**2 *  w**2  &
       - allrefs(ir,4 )**2 * kx**2  &
       + allrefs(ir,5 )*cmplx(0.,1.)*kx  &
       - allrefs(ir,6 )**2

       if ( aimag(carg) .ne. 0.) then !! Complex Check

          if ( real(carg) .lt. 0) then
             kz=-allrefs(ir,1)*kx+allrefs(ir,2)*cmplx(0.,1.)+csqrt(-carg)
             dax(ix)=dax(ix)*cexp(cmplx(-dz*abs(real(kz)),-dz*aimag(kz) ) )
          else
             kz=-allrefs(ir,1)*kx+allrefs(ir,2)*cmplx(0.,1.)+csqrt(+carg)
             dax(ix)=dax(ix)*cexp(cmplx(-dz*abs(aimag(kz)),-dz* real(kz) ))
          end if

       else
          if ( real(carg) .lt. 0) then
             kz =-dz*cmplx(allrefs(ir,2)+sqrt(-real(carg)),dir*(-kx*allrefs(ir,1)))
             dax(ix)=dax(ix)*exp(real(kz))*cmplx(cos(aimag(kz)),sin(aimag(kz)) )
          else
             kz =-dz*cmplx(allrefs(ir,2),dir*(sqrt(+real(carg))-kx*allrefs(ir,1)))
             dax(ix)=dax(ix)*exp(real(kz))*cmplx(cos(aimag(kz)),sin(aimag(kz)) )
          end if

       end if  !! End Complex check

    end do
    call fth( .true.,.false.,dax)

  end subroutine wemig_phs

  !----------------------------------------------------------------
  !! split-step Fourier
  subroutine wemig_ssf(dax,w,iz,ir,allfields,allrefs,dir) 
    complex :: dax(:)
    real    :: w,dir,allfields(:,:),allrefs(:,:)
    integer :: ir,iz

    ikz =(cmplx(0., 1.)*allrefs(ir,2)*( allfields(:,2 )-allrefs(ir,2 ) )) + &
               (   w**2*allrefs(ir,3)*( allfields(:,3 )-allrefs(ir,3 ) ))/ &
               sqrt( abs(allrefs(ir,3)**2*w**2 - allrefs(ir,6)**2)+eps ) 

    dax = dax * cexp( cmplx( -dz*abs(aimag(ikz)),-dz*dir*real(ikz)  ))

  end subroutine wemig_ssf
  !----------------------------------------------------------------
  subroutine rwetaper_init()
    integer :: itap,jtap,iy,ix
    real :: pi
    pi=acos(-1.)
    do itap=1,ntap
       jtap = abs(ntap-itap)
       tap(itap) = cos(pi/2.* jtap/(ntap-1))
       tap(nx-itap+1) = cos(pi/2.* jtap/(ntap-1))
    end do

  end subroutine rwetaper_init

  !----------------------------------------------------------------

  subroutine rwetaper(dax)
    complex, pointer :: dax(:)
    dax = dax * tap
  end subroutine rwetaper

  !----------------------------------------------------------------

end module  RWE2D_SI_wemig
