module RWE3D_SI_wemig
  use sep
  use util_
  use fft

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

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

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



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

contains
  !----------------------------------------------------------------  
  subroutine wemig_init(ax,ay,az,aw,ar,forward_in,kin_in,nref_in,verbose_in)
    type(myaxis)     :: ax,az,ay,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;    ny=ay%n;    nw=aw%n
    dz=az%d;    dx=ax%d;    dy=ay%d;    dw=aw%d * 2.*pi
    oz=az%o;    ox=ax%o;    oy=ay%o;    ow=aw%o * 2.*pi

    nref=nref_in

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

    dkx=2.*pi/(nx*dx)
    dky=2.*pi/(ny*dy)

    kin=kin_in
    verbose=verbose_in;    
    eps=0.0000001

    allocate(ikz2d(nx,ny))
    allocate(wfl(nx,ny))
    allocate(wtt(nx,ny))
    allocate(msk(nx,ny)); msk=1.
    allocate(mtt(nx,ny)); mtt=1.
    allocate(tap(nx,ny)); 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(:,:)

    !! . . Forward Scattering Option
  
    !! . . Begin Migration
    do iw=1,nw  !! Frequency Loop           
       w=ow+(iw-1)*dw
       rax => rwf(:,:,iw)
       do iz=2,nz  !! Depths
     
          wtt=rax;   rax=0. 
          do ir=1,nref  
             if (refs(ir,4,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*ny,8)

       end do !! . . z-loop
       write(0,*) maxval(abs(rax)),minval(abs(rax))
       if( mod(iw,2) .eq. 1) write(0,*) 'FINISHED FREQ ',iw,' OF ',nw 
    end do  !! . . w-loop
   call sep_write(rax,"SP")

  end subroutine wemig

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

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

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

    if (ny .gt. 1) then
       do ii=1,nloop !! Smoothing Iterations
          msk(1,:) = mtt(1,:)
          msk(:,1) = mtt(:,1)
          msk(nx,:)=mtt(nx-1,:)
          msk(:,ny)=mtt(:,ny-1)
          do ix=2,nx-1
             do iy=2,ny-1
                msk(ix,iy)=( mtt(ix-1,iy)+mtt(ix+1,iy)+mtt(ix,iy-1)+mtt(ix,iy+1) )/4.
             end do
          end do
          mtt=msk
       end do
    else
       do ii=1,nloop !! Smoothing Iterations
          msk(1,:) = mtt(1,:)
          msk(nx,:)=mtt(nx-1,:)
          do ix=2,nx-1
             msk(ix,1)=( mtt(ix-1,1)+mtt(ix+1,1) )/2.
          end do
          mtt=msk
       end do
    end if

    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(:,:)

    !! Initial Fourier Transform (x,y) -> (kx,ky)
    if (ny .eq. 1) then
       call ft1axis(.false.,.false.,dax)
    else
       call ft1axis(.false.,.false.,dax) ! Had it as .true.
       call ft2axis(.false.,.false.,dax)
    end if

    !! Loop over all kx and ky
    do ix=1,nx
       kx=okx+(ix-1)*dkx
       do iy=1,ny
          ky=oky+(iy-1)*dky

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

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

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

          else
 
              if ( real(carg) .lt. 0) then

                kz= -dz*cmplx( allrefs(ir,3) + sqrt( -real(carg) ),&
                           -kx*allrefs(ir,1) - ky*allrefs(ir,2) )

                dax(ix,iy)=dax(ix,iy)*exp(real(kz))*cmplx(cos(aimag(kz)),sin(aimag(kz)) )

             else
                !! ADDED EXTRA BRACKET
                kz =-dz* cmplx( allrefs(ir,3),&
                          + dir*sqrt(+real(carg)) - kx*allrefs(ir,1) - ky*allrefs(ir,2) )

                dax(ix,iy)=dax(ix,iy)*exp(real(kz))*cmplx(cos(aimag(kz)),sin(aimag(kz)) )

             end if

          end if  !! End Complex check
     
       end do
    end do

    !! Do Inverse Fourier transform (kx,ky) -> (x,y)
    if (ny .eq. 1) then
       call ft1axis( .true.,.false.,dax)
    else
       call ft2axis( .true.,.false.,dax)
       call ft1axis( .true.,.false.,dax) ! had .false.
    end if
  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

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

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

  end subroutine wemig_ssf
  !----------------------------------------------------------------
  subroutine rwetaper_init()
    integer :: itap,jtap,iy,ix
    real :: pi
    pi=acos(-1.)

    if (ny .eq. 1 ) then
       write(0,*) 'DOING 1D Taper'
       if (ntap .gt. 1) then
          !! Taper in x-direction
          do itap=1,ntap
             jtap = abs(ntap-itap)
              tap(   itap  ,1) = cos(pi/2.* jtap/(ntap-1))
          end do

          do itap=1,ntap
             jtap = abs(ntap-itap)
              tap(nx-itap+1,1) = cos(pi/2.* jtap/(ntap-1))
          end do

       end if
    else
       write(0,*) 'DOING 2D Taper'

       if(ntap>1) then
          !! Taper in x-direction
          do itap=1,ntap
             jtap = abs(ntap-itap)
             do iy=1,ny
                tap(   itap  ,iy) = tap(itap,iy) * cos(pi/2.* jtap/(ntap-1))
             end do
          end do

          do itap=1,ntap
             jtap = abs(ntap-itap)
             do iy=1,ny
                tap(nx-itap+1,iy) = tap(nx-itap+1,iy) * cos(pi/2.* jtap/(ntap-1))
             end do
          end do

          !! Taper in y-direction
          do itap=1,ntap
             jtap = abs(ntap-itap)
             do ix=1,nx
                tap(ix,itap  )    = tap(ix,itap) * cos(pi/2.* jtap/(ntap-1))
             end do
          end do

          do itap=1,ntap
             jtap = abs(ntap-itap)
             do ix=1,nx
                tap(ix,ny-itap+1) = tap(ix,ny-itap+1) * cos(pi/2.* jtap/(ntap-1))
             end do
          end do

       end if

    end if
  end subroutine rwetaper_init

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

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

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

end module  RWE3D_SI_wemig
