module RWE2D_TOPO_wemig
  !! Migration driver for the RWE2D migration program
  !! contains
  !!		1) wemig_init		- wemig initialization routine
  !!		2) wemig		- main migration driver routine
  !!		3) wemig_ssfboth	- split-step fourier propagation
  !!		4) wemig_phsboth	- phase-shift propagation
  !!		5) wemig_refboth	- interpolation for various phase-shifts
  !!		6) rwetaper_init	- a cosine-taper initialization routine
  !!		7) rwetaper		- apply cosine-taper to reduce boundary artifacts
  use sep		!! basic SEPlib routine I/O module
  use util		!! handle axes module
  use fft		!! FFT module - could use others though
  use Kernel_TOPO

  implicit none
  integer, private :: nt,nr,nw,ni,ntap
  real,    private :: dt,dr,dw,ot,or,ow,okr,dkr,eps,pi,sixth,okt,dkt

  !! PRIVATE ARRAYS FOR PROPAGATION
  integer, allocatable,private :: allmask(:)
  real,    allocatable,private :: tap(:)

  !! Global references
  real,    allocatable,private :: allflds(:,:,:)

  !! Propagation
  complex, allocatable,private :: swfl(:,:),swtt(:,:),rwfl(:,:),rwtt(:,:),rdax(:,:),sdax(:,:)

contains

  !----------------------------------------------------------------
  !! Initialization routine - parameter passing and array allocation
  subroutine wemig_init(ar,at,aw)
    type(myaxis)     :: ar,at,aw
    real             :: pi
    pi=acos(-1.)

    write(0,*) 'SETTING UP WEMIG'

    nt=at%n;    nr=ar%n;    nw=aw%n
    dt=at%d;    dr=ar%d;    dw=aw%d * 2.*pi
    ot=at%o;    or=ar%o;    ow=aw%o * 2.*pi

    call from_param("ntap",ntap,5)  	!! Width of cosine taper
    call from_param("sixth",sixth,0.122996)
    okr=  -pi/    dr
    dkr=2.*pi/(nr*dr)

    eps=0.0001
    allocate (tap(nr))
    tap=1.
    call rwetaper_init() !! Initialize the taper program
    write(0,*) 'SETUP WEMIG'

  end subroutine wemig_init

  !----------------------------------------------------------------
  !! WE migration driver
  !!
  !! Basic routine is:
  !!
  !!	for all frequencies
  !!		for all depths
  !!			Split-step Fourier Propagation
  !!			FFT wavefields
  !!			for all References
  !!				Phase-shift propagation
  !!				IFFT wavefields
  !! 				FD Correction
  !!				Wavefield Interpolation
  !!			end References
  !!			Imaging Condition
  !!		end depths
  !!	end frequencies
  !!
  subroutine wemig(rwf,swf,Pimg,iss,fields,vel,ddz)
    integer            :: iw,ir,it,ii,iss,id
    real               :: rarg,sarg,w
    real               :: Pimg(:,:),fields(:,:,:),vel(:,:),ddz(:,:)
    real,allocatable   :: vmin(:,:),veltmp(:,:,:,:),flds(:,:,:,:)
    real,allocatable   :: saxmax(:),Timg(:,:,:),ddd(:,:,:)
    complex            :: rwf(:,:),swf(:,:)
    complex,allocatable:: rax(:,:,:),sax(:,:,:)

    !! For OMP
    integer, external :: omp_get_num_threads
    integer, external :: omp_get_thread_num
    integer  :: nth,ith
    logical :: dofilt

    !$OMP PARALLEL
    nth = omp_get_num_threads()
    !$OMP END PARALLEL

    !! set and freeze the # of threads
    call omp_set_dynamic(.false.)
    call omp_set_num_threads(nth)

    allocate( sax(nr,1,nth), rax(nr,1,nth),veltmp(nr,1,nt,nth),saxmax(nth))
    allocate(Timg(nt,nr,nth),vmin(nt,nth),flds(nr,4,nt,nth),ddd(nr,nt,nth)  ) 

    !! . . Forward Scattering Option
    sarg=-1.; 
    rarg= 1.
    write(0,*) 'STARTING MIGRATION'

    !$OMP PARALLEL PRIVATE(id)
    id = omp_get_thread_num()
    write(0,*) 'HELLO WORLD from thread ', id
    !$OMP barrier
    if (id .eq. 0) then
       nth = omp_get_num_threads()
       write(0,*) 'THERE ARE ',nth,' threads'
    end if
    !$OMP END PARALLEL

    do it=1,nt
       do ith=1,nth
          vmin(it,ith) = minval(fields(:,2,it)*vel(it,:))
       end do
    end do

    do it=1,nt
       do ir=1,nr
          do ith=1,nth
             veltmp(ir,1,it,ith)=vel(it,ir)
          end do
       end do
    end do

    do it=1,nt
       do ir=1,nr
          do ii=1,4
             do ith=1,nth
                flds(ir,ii,it,ith)=fields(ir,ii,it)
             end do
          end do
       end do
    end do

    do it=1,nt
       do ir=1,nr
          do ith=1,nth
             ddd(ir,it,ith)=ddz(it,ir)
          end do
       end do
    end do

    !! CHECK HERE
    call shot_ker_init(nr, 1, dr , 1. ,dt , 0., 0., .11)

    !$OMP PARALLEL DO SCHEDULE(DYNAMIC) PRIVATE(iw,w,id,it,ir,tap) &
    SHARED(nw,ow,dw,nt,nr,Timg,sax,rax,flds,veltmp,vmin,swf,rwf,rarg,sarg,ddd)

    do iw=1,nw  !! FREQUENCY LOOP         

       id = omp_get_thread_num()+1

       w=ow+(iw-1)*dw 		
       saxmax(id)=maxval(cabs(swf(:,iw)))

       do ir=1,nr
          sax(ir,1,id) = swf(ir,iw)
       end do

       do ir=1,nr
          rax(ir,1,id) = rwf(ir,iw)
       end do

       call rwetaper(sax,id) 
       call rwetaper(rax,id)

       do it=2,nt  !! Time-step LOOP

          !$OMP CRITICAL
          call shot_ker_onestep(id,ddd(:,it,id),sax(:,:,id),veltmp(:,:,it,id),flds(:,:,it,id),w,sarg)
          !$OMP END CRITICAL

          !$OMP CRITICAL
          call shot_ker_onestep(id,ddd(:,it,id),rax(:,:,id),veltmp(:,:,it,id),flds(:,:,it,id),w,rarg)
          !$OMP END CRITICAL

          call fth(.false.,.false.,sax(:,1,id) )
          call fth(.false.,.false.,rax(:,1,id) )

          call wemig_phs_correction(id,ddd(:,it,id),sax(:,:,id),rax(:,:,id),w,vmin(it,id),sarg,rarg)

          call fth( .true.,.false.,sax(:,1,id) )
          call fth( .true.,.false.,rax(:,1,id) )

          call rwetaper(sax,id) 
          call rwetaper(rax,id) 

          do ir=1,nr
             Timg(it,ir,id)=Timg(it,ir,id) + &
                real(conjg(rax(ir,1,id))*sax(ir,1,id) )
               !! real(conjg(sax(ir,1,id))*sax(ir,1,id)+eps*saxmax(id) )
          end do
       end do !! End Depth

       if ( modulo( iw, nw/10) .eq. 0 ) &
       write(0,*) 'TIMG: ',ceiling(100.*dble(iw)/dble(nw)),nw,iss,maxval(Timg(:,:,id))

    end do  !! END Frequency Loop

    !$OMP END PARALLEL DO

    !! Imaging condition
    do it=1,nt
       do ir=1,nr
          do ith=1,nth
             Pimg(it,ir)=Pimg(it,ir)+Timg(it,ir,ith)
             Timg(it,ir,ith)=0.
          end do
       end do
    end do

    write(0,*) 'FINISHED OFF SHOT!'
    call shot_ker_release()
    deallocate( sax,rax,veltmp,vmin,saxmax,Timg,flds,ddd ) 

  end subroutine wemig


  !----------------------------------------------------------------
  !! Taper Initialization Program
  subroutine rwetaper_init()
    integer :: itap,jtap,ix
    real :: pi
    pi=acos(-1.)

    if (ntap .gt. 1) then

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

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

    end if

  end subroutine rwetaper_init

  !!----------------------------------------------------------------
  !!
  !! Apply Taper to wavefield
  !!
  subroutine rwetaper(dax,id)
    integer :: ir,id
    complex :: dax(:,:,:)
    do ir=1,nr
       dax(ir,1,id) = dax(ir,1,id) * tap(ir)
    end do
  end subroutine rwetaper


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

end module  RWE2D_TOPO_wemig
