module RWE3D_source_wemig
  use sep
  use util
  use fft
  use geometry

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

  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(Vel(nz,nx,ny))
    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,swf,Pimg,rays,Vel,ishot)
    complex, pointer :: rwf (:,:,:)
    complex, pointer :: swf (:,:,:)
    real             :: Pimg(:,:,:)
    real             :: rays(:,:,:,:)
    real             :: Vel(:,:,:)
    integer :: iw,iz,ix,iy,ir,ishot
    real    :: w,sarg,saxmax
    complex, pointer :: rax(:,:)
    complex, pointer :: sax(:,:)

    integer, allocatable :: mask(:,:,:)
    real,    allocatable :: refs(:,:,:), fields(:,:,:,:)

    allocate(mask(nx,ny,nz),refs(nref,10,nz),fields(nx,ny,10,nz) )
    mask=0; refs=0.;fields=0.;

    !! . . Forward Scattering Option
    sarg=-1.; if ( forward .ne. 0. ) sarg=1.

    !! . . Calculate All Geometric Aspects
    call geometry_init(nz,nx,ny,dz,dx,dy,kin,nref,verbose)
    call geometry_calc(rays,Vel,fields,refs,mask)
    call geometry_close()

    !! . . Begin Migration
    do iw=1,nw  !! Frequency Loop           
       w=ow+(iw-1)*dw
       sax => swf(:,:,iw)
       saxmax=maxval(abs(sax));
       do iz=2,nz 

          !! . . Do source wavefield
          wtt=sax;   sax=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),sarg)
             call wemig_ssf(wfl,w,iz,ir,fields(:,:,:,iz),refs(:,:,iz),sarg)       
             call wemig_ref(wfl,iz,ir,mask(:,:,iz),sax)

          end do
          call rwetaper(sax)

          Pimg(iz,:,:) = Pimg(iz,:,:) + real( sax )


       end do !! . . z-loop

       write(0,*) "SSF",iw,nw,ishot,maxval(abs(Pimg(:,:,:)))

    end do  !! . . w-loop

  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 = -dz*cmplx(  -abs(allrefs(ir,3)) + sqrt(-abs(real(carg)))  , & 
                       dir*( - kx*allrefs(ir,1)- ky*allrefs(ir,2)+sqrt( abs(aimag(carg)))) )

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

             else

                kz = -dz*cmplx(  -abs(allrefs(ir,3)) + sqrt( abs(aimag(carg)))  ,&
                          + dir*(sqrt( abs(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

          else

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

                kz =-dz*cmplx( -abs(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
                
                kz =-dz*cmplx( -abs(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 ) )) + &
    ikz2d =(cmplx(0., 0.)*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_source_wemig
