module RWE2D_GOGI_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,ntap
  real,    private :: dr,dt,dw,or,ot,ow,okr,dkr,forward,pi,okt,dkt,r0,rend,rbeg,eps
  logical, private :: kin,verbose

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

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

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

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

contains

  !----------------------------------------------------------------
  !! Initialization routine - parameter passing and array allocation
  subroutine wemig_init(ar,at,aw,forward_in,kin_in,verbose_in,r0_in,rbeg_in,rend_in,eps_in)
    type(myaxis)     :: ar,at,aw
    real             ::          forward_in,                          r0_in,rbeg_in,rend_in,eps_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,ir

    rbeg = rbeg_in
    rend = rend_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;
    eps=eps_in
    dw=aw%d * 2.*pi; ow=aw%o * 2.*pi
    verbose=verbose_in;    
    r0=r0_in

    !! Geometry
    allocate(rline(nr),tline(nt))
    do ir = 1,nr
       rline(ir) = acos( -1. + (rbeg + (ir-1)*ar%d)/ r0  )
    end do

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

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

    kin=kin_in

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

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

    !! Propagation
    allocate( rdax(nr,nth), sdax(nr,nth), rwfl(nr,nth), swfl(nr,nth), rwtt(nr,nth), swtt(nr,nth),  tap(nr) )

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

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

    if (verbose) then
       call swf_print()
       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,swf,Pimg,ishot,vel)
    integer            :: iw,ir,it,inn,is,id,ip,ishot
    real               :: rarg,sarg,ss,w,Pimg(:,:),vel(:,:)
    complex            :: rwf(:,:),swf(:,:)
    complex,allocatable:: rax(:,:,:),sax(:,:,:) 
    real   ,allocatable:: veltmp(:,:,:),vmin(:),saxmax(:)
    !! 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( sax(nr,1,nth), rax(nr,1,nth), veltmp(nr,1,nt),vmin(nt),saxmax(nth)  ) 

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

    !! . . Forward Scattering Option
    sarg=-1.; !! if ( forward .ne. 0. ) sarg=1.
    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,*) "DR and DT: ",dr,r0*dt

    call shot_ker_init(nr, 1, 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,rarg,sarg,Timg,sax,rax,allflds,saxmax)

    do iw=1,nw  !! FREQUENCY LOOP         
       id = omp_get_thread_num()+1

       w = (iw-1)*dw+ow

       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)

       saxmax(id) = maxval( real(conjg(sax(:,1,id))*sax(:,1,id)))

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

          !! Step Source wavefield
          !$OMP CRITICAL
          call shot_ker_onestep(id,sax,veltmp(:,:,it),allflds(:,:,it),w,sarg)
          !$OMP END CRITICAL

          !! 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.,sax(:,1,id) )
          !$OMP END CRITICAL

          !$OMP CRITICAL
          call fth(.false.,.false.,rax(:,1,id) )
          !$OMP END CRITICAL

          !! Phase correction routine
          !$OMP CRITICAL
          call wemig_phs_correction(id,sax(:,1,:),rax(:,1,:),w,vmin(it),sarg,rarg)
          !$OMP END CRITICAL

          !! FFT to x
          !$OMP CRITICAL
          call fth( .true.,.false.,sax(:,1,id) )
          !$OMP END CRITICAL

          !$OMP CRITICAL
          call fth( .true.,.false.,rax(:,1,id) )
          !$OMP END CRITICAL

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

          !! Imaging condition
          !$OMP CRITICAL
          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 )
          end do
          !$OMP END CRITICAL

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

    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( sax , rax , veltmp,  vmin , saxmax ) 

  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 swf_print()
    call auxputch("n1","i",nr,"swf.H")
    call auxputch("n2","i",nt,"swf.H")
    call auxputch("d1","r",dr,"swf.H")
    call auxputch("d2","r",dt,"swf.H")
    call auxputch("o1","r",or,"swf.H")
    call auxputch("o2","r",ot,"swf.H")
    call auxputch("esize","i",8,"swf.H")
    call auxputch("n3","i",nw,"swf.H")
    call auxputch("d3","r",dw,"swf.H")
    call auxputch("o3","r",ow,"swf.H")
  end subroutine swf_print

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

  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_GOGI_wemig
