module RWE2D_ZERO_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_gogi    !! GOGI's 80 degree code

  implicit none
  integer, private :: nr,nt,nw,nloop,nref,ntap
  real,    private :: dr,dt,dw,or,ot,ow,okr,dkr,forward,eps,pi,okt,dkt,r0,maxang,minang
  logical, private :: kin,verbose

  !! Geometry
  real,    allocatable,private :: rline(:),tline(:)

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


  !! Temp in/out to set up global refs
  real,    allocatable,private :: refs(:,:),flds(:,:)

  !! Propagation
  real,    allocatable,private :: msk(:),mtt(:),tap(:)
  complex, allocatable,private :: rwfl(:,:),rwtt(:,:),rdax(:,:)

  !! Temp. image
  real,    allocatable,private :: Timg(:,:,:)

contains

  !----------------------------------------------------------------
  !! Initialization routine - parameter passing and array allocation
  subroutine wemig_init(ar,at,aw,forward_in,kin_in,nref_in,verbose_in,r0_in,minang_in,maxang_in)
    type(myaxis)     :: ar,at,aw
    integer          ::                            nref_in
    real             ::          forward_in,                          r0_in,minang_in,maxang_in
    logical          ::                     kin_in,        verbose_in
    real             :: pi

    !! For OMP
    integer, external :: omp_get_num_threads
    integer, external :: omp_get_thread_num
    integer  :: nth,ith,i1

    minang = minang_in
    maxang = maxang_in

    pi=acos(-1.)
    forward=forward_in
    nr=ar%n;    nt=at%n;
    dr=ar%d;    dt=at%d; 
    or=ar%o;    ot=at%o; 
    nw=aw%n;

    dw=aw%d * 2.*pi; ow=aw%o * 2.*pi
    verbose=verbose_in;    
    r0=r0_in
    nref=nref_in

    !! Geometry
    allocate(rline(nr),tline(nt))

    !! Receiver line
    rline=(/ (i1,i1=0,nr-1) /)*dr+or
    tline=(/ (i1,i1=0,nt-1) /)*dt+ot


    call from_param("ntap" ,ntap, 5) !! Width of cosine taper
    call from_param("nloop",nloop,3) !! Number of smoothing iterations 
    okr= -pi/dr;  dkr=2.*pi/(nr*dr)
    okt= -pi/dt;  dkt=2.*pi/(nt*dt)

    kin=kin_in
    eps=0.0000001

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

    !! Global Reference allocation
    allocate( allflds(nr,1  ,nt) ) 
    allocate( allrefs(nref,1,nt) )

    !! Temp in/out to set up global refs
    allocate( flds(nr,1) , refs(nref,1) )


    !! Propagation
    allocate( rdax(nr,nth), rwfl(nr,nth), rwtt(nr,nth) )
    allocate( msk(nr), mtt(nr), tap(nr) )

    !! Temp Image
    allocate( Timg(nt,nr,nth) )
    Timg=0.

    tap=1.;    msk=0.; 
    call rwetaper_init() !! Initialize the taper program

    if (verbose) then
       call rwf_print()
       call img_print()
    end if

    !! Set up GOGI's Propagators
  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
  !!				Wavefield Interpolation
  !!			end References
  !!			Imaging Condition
  !!		end depths
  !!	end frequencies
  !!
  subroutine wemig(rwf,Pimg,ishot,vel)
    integer            :: iw,ir,it,inn,is,id,ip,ishot
    real               :: rarg,ss,w,Pimg(:,:),vel(:,:)
    complex            :: rwf(:,:)
    complex,allocatable:: rax(:,:,:)
    real   ,allocatable:: veltmp(:,:,:),vmin(:)
    !! 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

    allocate(  rax(nr,1,nth), veltmp(nr,2,nt), vmin(nt)  ) 

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

    !! . . Forward Scattering Option
    rarg=1.

    !$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


    !! Geometry
    do it=1,nt
       do ir=1,nr  
          allflds(ir,1,it) = sqrt( sinh(tline(it))**2 + sin(rline(ir))**2)
          veltmp (ir,1,it) = vel(it,ir)
       end do
       vmin(it) = minval(veltmp(:,1,it)/allflds(:,1,it))
    end do

    write(0,*) 'STARTING MIGRATION'
    call shot_ker_init(nr, 1, r0*dr  , 1., r0*dt, 0., 0., .11)

    !$OMP PARALLEL DO SCHEDULE(DYNAMIC) PRIVATE(iw,w,id,it,ir,tap) SHARED(nw,ow,dw,nt,nr,nref,rarg,Timg,rax,allflds)
    
    do iw=1,nw  !! FREQUENCY LOOP         
       id = omp_get_thread_num()+1

       w = (iw-1)*dw+ow
   
       do ir=1,nr
          rax(ir,1,id) = rwf(ir,iw)
       end do

       call rwetaper(rax,id)

       do it=2,nt  !! Time-step LOOP
     
          !! Step Receiver wavefield
          !$OMP CRITICAL
          call shot_ker_onestep(id,rax,veltmp(:,:,it),allflds(:,:,it),w,rarg)
          !$OMP END CRITICAL

          !! High-angle filter 
          !! FFT to kx
          !$OMP CRITICAL
          call fth(.false.,.false.,rax(:,1,id) )
          !$OMP END CRITICAL

          !! Phase correction routine
          !$OMP CRITICAL
          call correction(id,rax(:,1,:),w,vmin(it),rarg)
          !$OMP END CRITICAL
          
          !! FFT to x
          !$OMP CRITICAL
          call fth( .true.,.false.,rax(:,1,id) )
          !$OMP END CRITICAL

          !! Taper
          call rwetaper(rax,id) 

          !! Imaging condition
          !$OMP CRITICAL
          do ir=1,nr
             Timg(it,ir,id)=Timg(it,ir,id)+ real( rax(ir,1,id) )
          end do
          !$OMP END CRITICAL

       end do !! End Depth
       if ( modulo( iw, max(nw/10,1)) .eq. 0 ) then
          write(0,*) 'TIMG: ',ceiling(100.*dble(iw)/dble(nw)),ishot,maxval(Timg(:,:,id))
       end if

    end do !! END frequency loop

    !$OMP END PARALLEL DO

    !! SUMMATION OVER OUTSIDE IMAGE CONDITION
    do ir=1,nr
       do it=1,nt
          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

    !! Deallocate
    call shot_ker_release()
    deallocate( rax , veltmp,  vmin  ) 

  end subroutine wemig

  !!
  !----------------------------------------------------------------
  !! Taper Initialization Program
  subroutine rwetaper_init()
    integer :: itap,jtap
    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)
    complex  :: dax(:,:,:)
    integer :: ir,id
    do ir=1,nr
       dax(ir,1,id) = dax(ir,1,id) * tap(ir)
    end do
  end subroutine rwetaper

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

  subroutine rwf_print()
    call auxputch("n1","i",nr,"rwf.H")
    call auxputch("n2","i",nt,"rwf.H")
    call auxputch("d1","r",dr,"rwf.H")
    call auxputch("d2","r",dt,"rwf.H")
    call auxputch("o1","r",or,"rwf.H")
    call auxputch("o2","r",ot,"rwf.H")
    call auxputch("esize","i",8,"rwf.H")
    call auxputch("n3","i",nw,"rwf.H")
    call auxputch("d3","r",dw,"rwf.H")
    call auxputch("o3","r",ow,"rwf.H")
  end subroutine rwf_print
  !----------------------------------------------------------------

  subroutine img_print()
    call auxputch("n1","i",nr,"img.H")
    call auxputch("n2","i",nt,"img.H")
    call auxputch("d1","r",dr,"img.H")
    call auxputch("d2","r",dt,"img.H")
    call auxputch("o1","r",or,"img.H")
    call auxputch("o2","r",ot,"img.H")
    call auxputch("esize","i",4,"img.H")
    call auxputch("n3","i",nw,"img.H")
    call auxputch("d3","r",dw,"img.H")
    call auxputch("o3","r",ow,"img.H")
  end subroutine img_print


 
  !----------------------------------------------------------------
end module  RWE2D_ZERO_wemig
