module RWE3D_source_wemig
  use sep
  use util_
  use fft
  use geometry
  use sep_timers_mod

  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,ikz2d
  complex, allocatable,dimension(:,:),private :: swfl,swtt,sdax

  real,    allocatable,dimension(:,:  ),private :: allrefs
  integer, allocatable,dimension(:,:  ),private :: allmask
  real,    allocatable,dimension(:,:,:),private :: allfields

contains
  !----------------------------------------------------------------  
  subroutine wemig_src_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

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

    kin=kin_in
    verbose=verbose_in;    
    eps=0.0000001

    allocate(sdax(nx,ny))
    allocate(ikz2d(nx,ny))
    allocate(swfl(nx,ny))
    allocate(swtt(nx,ny))
    allocate(msk(nx,ny))
    allocate(mtt(nx,ny))
    allocate(tap(nx,ny))
    allocate(allmask(nx,ny))
    allocate(allfields(nx,ny,10))
    allocate(allrefs(nref,10))

    msk=1;mtt=1.;tap=1.
    call rwetaper_init()

  end subroutine wemig_src_init
 
  !----------------------------------------------------------------

  !! WE migration driver
  subroutine wemig_src(swf,Pimg,iss,fields,refs,rmask)
    integer          :: iw,iz,ix,iy,ir,iss,rmask(:,:,:)
    real             :: Pimg(:,:,:)
    real             :: w,sarg,refs(:,:,:),fields(:,:,:,:)
    complex, pointer :: swf(:,:,:),sax(:,:)
    logical          :: logic
    integer          :: t1,t2,t3,t4,t5,t6

    logic=init_sep_timers()
    logic=setup_next_timer("WLoop",t1)
    logic=setup_next_timer("PHS",t2)
    logic=setup_next_timer("SSF",t3)
    logic=setup_next_timer("REF",t4)

    !! . . Forward Scattering Option
    sarg=1.; !! if ( forward .ne. 0. ) sarg=1.
    write(0,*) 'STARTING MIGRATION'
    write(0,*) 'USING REFERENCE MAX OF :',nref

    do iw=1,nw  !! Frequency Loop           
       call start_timer_num(t1)

       w=ow+(iw-1)*dw
       sax => swf(:,:,iw)
       do iz=2,nz

          call start_timer_num(t3)
          call wemig_ssfboth(sax,w,rmask(:,:,iz),fields(:,:,:,iz),refs(:,:,iz),sarg)
          call stop_timer_num(t3)
          
          swtt=sax;sax=0.

          call ft1axis(.false.,.false.,swtt) 
          call ft2axis(.false.,.false.,swtt)

	  do ir=1,nref  
             if (refs(ir,4,iz) .lt. 0.000000001) cycle 
             swfl=swtt
!             write(0,*) iw,iz,ir

!            write(0,*) 'SWFL B',minval(abs(swfl)),maxval(abs(swfl))

             call start_timer_num(t2)
             call wemig_phsboth(swfl,w,iz,ir,refs(:,:,iz),sarg)
             call stop_timer_num(t2)

!             write(0,*) 'SWFL A',minval(abs(swfl)),maxval(abs(swfl))

             call ft2axis( .true.,.false.,swfl)
             call ft1axis( .true.,.false.,swfl) 

 !            write(0,*) 'SWFL FFT',minval(abs(swfl)),maxval(abs(swfl))

             call start_timer_num(t4)
             call wemig_refboth(swfl,iz,ir,rmask(:,:,iz),sax)
             call stop_timer_num(t4)

          end do

          call rwetaper(sax)
  !        write(0,*) 'SAX',iw,iz,minval(abs(sax)),maxval(abs(sax))

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

       end do !! . . z-loop

       write(0,*) "SSF",iw,nw,iss,maxval(abs(Pimg(:,:,:)))
       call stop_timer_num(t1)

    end do  !! . . w-loop
    write(0,*) 'FINISHED OFF SHOT!'

    call print_timers()
  end subroutine wemig_src

  !----------------------------------------------------------------
  !! split-step Fourier
  subroutine wemig_ssfboth(sdax,w,allmask,allfields,allrefs,sdir) 
    integer :: ix,iy,allmask(:,:)
    real    :: w,sdir,rdir,allfields(:,:,:),allrefs(:,:)
    complex :: sdax(:,:)
    
    do iy=1,ny
       do ix=1,nx
          ikz2d(ix,iy)=-dz*(w**2*allrefs(allmask(ix,iy),4 )*&
            (allfields(ix,iy,4 )-allrefs(allmask(ix,iy),4))-&
            allrefs(allmask(ix,iy),10)*(allfields(ix,iy,10)-allrefs(allmask(ix,iy),10)))/&
            sqrt(abs(allrefs(allmask(ix,iy),4)**2*w**2-allrefs(allmask(ix,iy),10)**2)+eps ) 
       end do
    end do

    sdax = sdax*cmplx( cos(ikz2d) ,sdir*sin(ikz2d) )
 
  end subroutine wemig_ssfboth
  
  !----------------------------------------------------------------  

  !! Fourier-domain phase shift
  subroutine wemig_phsboth(sdax,w,iz,ir,allrefs,sdir)
   integer :: ix,iy,iz,ir
   real    :: w,kx,ky,sdir,carg,allrefs(:,:),pkz
   complex :: kz,sdax(:,:)


   do iy=1,ny
      ky=oky+(iy-1)*dky
      do ix=1,nx
         kx=okx+(ix-1)*dkx

         carg = &
           (allrefs(ir,4)* w)**2  &
         - (allrefs(ir,5)*kx)**2  &
         - (allrefs(ir,6)*ky)**2  &
         -  allrefs(ir,7)*kx*ky  &
         -  allrefs(ir,10)**2

         if ( carg .lt. 0) then

            pkz = -abs(dz*sqrt(-carg))
            sdax(ix,iy)=sdax(ix,iy)*exp(pkz)   

         else

            pkz =-dz*(sqrt(carg)-kx*allrefs(ir,1)-ky*allrefs(ir,2))

            sdax(ix,iy)=sdax(ix,iy)*cmplx(cos(pkz),sdir*sin(pkz))

         end if
      end do
   end do
 
  end subroutine wemig_phsboth

 
   !----------------------------------------------------------------
  subroutine wemig_refboth(swfl,iz,ir,allmask,sdax)
    complex :: sdax(:,:),swfl(:,:)
    integer :: iz,ir,ix,iy,ii,allmask(:,:)

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

    do ii=1,nloop !! Smoothing Iterations
       msk(1 ,1) = mtt(1   ,1)
       msk(nx,1) = mtt(nx-1,1)
       msk(1,ny) = mtt(1,ny-1)
       msk(nx,ny)= mtt(nx ,ny)
       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
    sdax = sdax + msk*swfl

 end subroutine wemig_refboth

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

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

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

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

       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
    write(0,*) 'TAPER',minval(tap),maxval(tap)
  end subroutine rwetaper_init

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

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


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

end module  RWE3D_source_wemig
