module RWE2D_ELLIPTIC_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 lloyd_mod !! Lloyd's code for velocity picking

  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
  integer, allocatable,private :: allmask(:,:  )
  real,    allocatable,private :: allrefs(:,:,:),allflds(:,:,:)


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

  !! Inside lloyds
  integer, allocatable,private :: iregion(:),nstart(:)
  real,    allocatable,private :: slice(:,:),scaly(:)

  !! Propagation
  real,    allocatable,private :: msk(:),mtt(:),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,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( allmask(nr,    nt) )
    allocate( allflds(nr,2  ,nt) ) 
    allocate( allrefs(nref,2,nt) )

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

    !! Inside lloyds
    allocate( scaly(2) , iregion(nr) ,  nstart(2) , slice(nr,2) )

    !! Propagation
    allocate( rdax(nr,nth), sdax(nr,nth), rwfl(nr,nth), swfl(nr,nth), rwtt(nr,nth), swtt(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 swf_print()
       call rwf_print()
       call img_print()
       call ref1_print()
       call ref2_print()
       call a_print()
    end if

    !! Calculate Geometry for swf and rwf
    call init_lloyd()

  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,saxmax,sarg,ss,w,Pimg(:,:),vel(:,:)
    complex            :: rwf(:,:),swf(:,:)
    complex,allocatable:: rax(:,:),sax(:,:) 

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

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


    allocate( sax(nr,nth), rax(nr,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  

          flds(ir,1)=abs(r0)*sqrt( sinh(tline(it))**2 + sin(rline(ir))**2 )/ vel(it,ir)
          flds(ir,2)=abs(r0)*sqrt( sinh(tline(it))**2 + sin(rline(ir))**2 )/ vel(it,ir)
       end do

       !if (modulo(it,nt/10) .eq. 0) write(0,*) "Ref",it,nt
       call get_references(flds,refs,mask)

       do ir=1,nr
          allflds(ir,1,it) = flds(ir,1)
          allflds(ir,2,it) = flds(ir,2)
          allmask(ir,  it) = mask(ir  ) 
       end do

       do is=1,nref
          allrefs(is,1,it)=refs(is,1)
          allrefs(is,2,it)=refs(is,2)
       end do
       refs=0. ; flds=0.; mask=1
    end do

    write(0,*) 'STARTING MIGRATION'

    !$OMP PARALLEL DO SCHEDULE(DYNAMIC) PRIVATE(iw,w,id,it,ir,inn,tap) &
    SHARED(nw,ow,dw,nt,nr,nref,rarg,sarg,Timg,sax,rax,swtt,rwtt,swfl,rwfl,allflds,allmask,allrefs)

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

       w = (iw-1)*dw+ow

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

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

       ! saxmax=maxval(cabs(sax(:,id)))

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

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

          call wemig_ssf(id,sax,rax,w,allmask(:,it),allflds(:,:,it),allrefs(:,:,it),sarg,rarg)

          do ir=1,nr
             swtt(ir,id)=sax(ir,id);
             sax(ir,id)=0.
          end do

          do ir=1,nr
             rwtt(ir,id)=rax(ir,id);
             rax(ir,id)=0.
          end do

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

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

          do inn=1,nref
             if ( allrefs(inn,1,it) .eq. -1.) cycle

             do ir=1,nr
                swfl(ir,id)=swtt(ir,id)
             end do

             do ir=1,nr
                rwfl(ir,id)=rwtt(ir,id)
             end do

             call wemig_phs(id,swfl,rwfl,w,inn,allrefs(:,:,it),sarg,rarg,it)

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

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

             call wemig_ref(id,swfl,rwfl,inn,it,allmask(:,it),sax,rax)
 
         end do


          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,id))*sax(ir,id))
!!                  real( cmplx(cos(-1.0*w),sin(-1.0*w)) * sax(ir,id) )
         end do
          !$OMP END CRITICAL

       end do !! End Depth

       if ( id .eq. 1) write(0,*) 'MAXVAL TIMG: ',iw,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


  end subroutine wemig

  !----------------------------------------------------------------
  !!
  !! Split-step Fourier Propagation subroutine 
  !!	
  subroutine wemig_ssf(id,sdax,rdax,w,mask,flds,refs,sarg,rarg) 
    integer :: ir,mask(:),id
    real    :: w,sarg,rarg,flds(:,:),refs(:,:),iikt2d,tmp
    complex :: sdax(:,:),rdax(:,:)

    !! Calculate w-r domain correction step
    do ir=1,nr
       iikt2d = -dt*w*( flds(ir,1) - refs(mask(ir),1))
       sdax(ir,id) =  sdax(ir,id) * cmplx( cos(iikt2d) , sarg*sin(iikt2d) ) 
    end do

    do ir=1,nr
       iikt2d = -dt*w*( flds(ir,1) - refs(mask(ir),1))
       rdax(ir,id) =  rdax(ir,id) * cmplx( cos(iikt2d) , rarg*sin(iikt2d) )  
    end do

  end subroutine wemig_ssf

 
  !----------------------------------------------------------------  
  !!
  !! Fourier-domain phase shift
  !!
  subroutine wemig_phs(id,sdax,rdax,w,inn,refs,sarg,rarg,it)
    integer :: ir,inn,it,id
    real    :: w,kr,sarg,rarg,carg,carg1,refs(:,:),pkr
    complex :: sdax(:,:),rdax(:,:)

    carg1= (w*refs(inn,1))**2
    do ir=1,nr 
       kr=okr+(ir-1)*dkr  
       carg = carg1 - kr**2

       if ( carg .lt. 0) then 

          pkr = -dt*abs(sqrt(-carg)) 
          
          sdax(ir,id) = sdax(ir,id)*exp(pkr) 
          rdax(ir,id) = rdax(ir,id)*exp(pkr) 

       else 

          pkr = -dt*sqrt(carg) 
          sdax(ir,id) = sdax(ir,id) * cmplx(cos(pkr),sarg*sin(pkr)) 
          rdax(ir,id) = rdax(ir,id) * cmplx(cos(pkr),rarg*sin(pkr)) 
       end if
    end do  

  end subroutine wemig_phs

 
  !!----------------------------------------------------------------
  !!
  !! Interpolation routine for different phase-shifts
  !!
  subroutine wemig_ref(id,swfl,rwfl,inn,it,mask,sdax,rdax)
    integer :: mask(:),id
    complex :: sdax(:,:),rdax(:,:),swfl(:,:),rwfl(:,:)
    integer :: ir,inn,ii,it

    !! Find out locations where want to use wavefield for given
    !! Reference values
    mtt=0.;
    where( mask .eq. inn ) 
       mtt=1.
    end where

    do ii=1,nloop
       msk(1 ) = mtt(1   )
       msk(nr) = mtt(nr-1)
       do ir=2,nr-1
          msk(ir)=( mtt(ir-1)+mtt(ir+1) )/2.
       end do
       do ir=1,nr
          mtt(ir)=msk(ir)
       end do
    end do

    do ir=1,nr
       sdax(ir,id) = sdax(ir,id) + msk(ir)*swfl(ir,id)
    end do
    
    do ir=1,nr
       rdax(ir,id) = rdax(ir,id) + msk(ir)*rwfl(ir,id)
    end do
  end subroutine wemig_ref


  !----------------------------------------------------------------
  !! 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 :: ip,id
    do ip=1,nr
       dax(ip,id) = dax(ip,id) * tap(ip)
    end do
  end subroutine rwetaper

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

  subroutine get_references(flds,refs,mask)
    integer ::iref,ii,ix,ir,id,icnt
    real   :: s1,s2
    real   :: flds(:,:)
    real   :: refs(:,:)
    integer:: mask(:)

    nstart = (/ 10,10 /)

    s1 = sum(flds(:,1))/dble(nr) 
    s2 = sum(flds(:,2))/dble(nr) 

    do ir=1,nr
       slice(ir,1)=flds(ir,1)/s1
       slice(ir,2)=flds(ir,2)/s2
    end do

  
    iref=lloyd_go(slice,refs,nref,iregion,nstart) 

    do ir=1,nref
       refs(ir,1) = refs(ir,1)*s1
       refs(ir,2) = refs(ir,2)*s2
    end do

    if(iref < nref) refs(iref+1:nref,:)=-1.

    do ir=1,nr
       mask(ir)=iregion(ir) 
    end do

    if (verbose) then
       call srite("ref1.H" ,refs(mask,1),size(flds(:,1))*4)
       call srite("ref2.H" ,refs(mask,2),size(flds(:,2))*4)
       call srite("a1.H" ,  flds(:,1),size(flds(:,1))*4)
       call srite("a2.H" ,  flds(:,2),size(flds(:,2))*4)
    end if

  end subroutine get_references

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

  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

  subroutine ref1_print()
    call auxputch("n1","i",nr  ,"ref1.H")
    call auxputch("n2","i",nt  ,"ref1.H")
    call auxputch("d1","r",dr  ,"ref1.H")
    call auxputch("d2","r",dt  ,"ref1.H")
    call auxputch("o1","r",or  ,"ref1.H")
    call auxputch("o2","r",ot  ,"ref1.H")
    call auxputch("esize","i",4,"ref1.H")
  end subroutine ref1_print

  subroutine ref2_print()
    call auxputch("n1","i",nr  ,"ref2.H")
    call auxputch("n2","i",nt  ,"ref2.H")
    call auxputch("d1","r",dr  ,"ref2.H")
    call auxputch("d2","r",dt  ,"ref2.H")
    call auxputch("o1","r",or  ,"ref2.H")
    call auxputch("o2","r",ot  ,"ref2.H")
    call auxputch("esize","i",4,"ref2.H")
  end subroutine ref2_print

 
  subroutine a_print()
    call auxputch("n1","i",nr,"a1.H" )
    call auxputch("n1","i",nr,"a2.H" )
    call auxputch("n2","i",nt,"a1.H" )
    call auxputch("n2","i",nt,"a2.H" )
    call auxputch("o1","r",0.,"a1.H" )
    call auxputch("o1","r",0.,"a2.H" )
    call auxputch("o2","r",0.,"a1.H" )
    call auxputch("o2","r",0.,"a2.H" )
    call auxputch("d1","r",1.,"a1.H" )
    call auxputch("d1","r",1.,"a2.H" )
    call auxputch("d2","r",1.,"a1.H" )
    call auxputch("d2","r",1.,"a2.H" )
  end subroutine a_print
 
  !----------------------------------------------------------------
end module  RWE2D_ELLIPTIC_wemig
