module RWE2D_source_wemig
  use sep
  use util_
  use fft
  use Kernel_gogi
  use C2R2D
  implicit none
  integer, private :: nz,nx,nw,nr,nt,rrn,i1,ir,it
  integer, private :: nth,ntap,nsx,nsz,norm,ns
  real,    private :: dz,dx,dw,oz,ox,ow,pi,dt,ot,rro,rrd,or,dr,r0
  real,    private :: ds,os
  integer, allocatable,private :: near_rec(:)
  real,    allocatable,private :: rline(:),tline(:)
  real,    allocatable,private :: flds(:,:,:)
  real,    allocatable,private :: tap(:),wflds(:,:)
  real,    allocatable,private :: Timg(:,:,:),gnorm(:,:),CARTvel(:,:)
  complex, allocatable,private :: sax(:,:)
  real,    allocatable,private :: raysin(:,:,:),vmin(:),RWEvel(:,:),veltmp(:,:,:)

contains

  !----------------------------------------------------------------  
  subroutine wemig_init(ax,az,aw,as,at,ar,rr,Vel)
    type(myaxis)     :: ax,az,aw,as,at,ar,rr
    real             :: pi,x
    real             :: Vel(:,:)
    integer          :: ix,iz,it,ir,iis(1)

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

    pi=acos(-1.)
    nz=az%n;    nx=ax%n;    nw=aw%n
    dz=az%d;    dx=ax%d;    dw=aw%d * 2.*pi
    oz=az%o;    ox=ax%o;    ow=aw%o * 2.*pi
    nr=ar%n;    nt=at%n;    rrn=rr%n
    or=ar%o;    ot=at%o;    rro=rr%o
    dr=ar%d;    dt=at%d;    rrd=rr%d
    ns=as%n;    ds=as%d;    os=as%o
    norm=1;
  
    !$OMP PARALLEL
    nth = omp_get_num_threads()
    !$OMP END PARALLEL

    allocate(sax(nt,1),raysin(nr,nt,2),gnorm(nr,nt),vmin(nr) )
    allocate(flds(nt,1,nr),RWEvel(nr,nt),veltmp(nt,1,nr),CARTvel(nz,nx)  )
    allocate(tap(nt),near_rec(nr)  )
    allocate(rline(nr),tline(nt),wflds(nr,nt) )

    call from_param("nsx",nsx,2)
    call from_param("nsz",nsz,2)
    call from_param("ntap",ntap,5)
    call from_param("norm",norm,1)

    do it=1,nt
       tap(it)=1.
    end do
    call rwetaper_init()
    
    r0 = ( (rr%n-1.)*rr%d )/2.
    write(0,*) "Old r0 value is; ",r0,nr,or,dr

    !! Make rays for interpolation
    rline=(/ (i1,i1=0,nr-1) /)* ar%d + ar%o
    tline=(/ (i1,i1=0,nt-1) /)* abs(2.*ot)/dble(nt) + ot

    !! GENERATE COORDINATE SYSTEM
    do it=1,nt
       do ir=1,nr
          raysin(ir,it,1) = -abs(r0)*cosh(tline(it))*   cos(rline(ir)    )+(r0+rro)
          raysin(ir,it,2) =  abs(r0)*sinh(tline(it))*   sin(rline(ir)    )
          flds  (it,1,ir) =    sqrt( sinh(tline(it))**2+sin(rline(ir))**2)
          gnorm (ir,it  ) =  r0*r0*( sinh(tline(it))**2+sin(rline(ir))**2) 
          wflds (ir,it  ) =    sqrt( sinh(tline(it))**2+sin(rline(ir))**2)
       end do
    end do
    call sep_write(cmplx(raysin(:,:,1),raysin(:,:,2)),"rays.H")
    call sep_write(wflds,"velocity.H")
    call coord_init( -raysin(1,nt/2,1) ) 
    
    do ir=1,nr
       x = (ir-1)*rrd+rro
       iis =  minloc( abs(x - raysin(:,nt/2,1)) )
       near_rec(ir) = iis(1)
    end do

    call coord_init( as%o ) 

    do ix=1,nx
       do iz=1,nz
          CARTvel(iz,ix)=Vel(iz,ix)
       end do
    end do

  end subroutine wemig_init

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

  subroutine coord_init( shot_shift )
    real :: shot_shift
   
    do it=1,nt
       do ir=1,nr
          raysin(ir,it,1) = raysin(ir,it,1)+shot_shift
       end do
    end do

  end subroutine coord_init

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

  !! WE migration driver
  subroutine wemig(adj,add,swf,swfout)
    integer            :: iw,iz,ix,ir,iss,id
    real               :: sarg,w,dd1,dd2,dx1,dx2,x
    complex            :: swf(:,:),swfout(:,:)
    complex,allocatable:: Pimg(:,:)
    logical            :: Iadj,adj,add


    allocate( Pimg(nt,nr) )

    write(0,*) 'using: ',r0*dt,r0*dr
    call shot_ker_init(nt, 1, r0*dt  , 1., r0*dr, 0., 0., .11)

    !! Get RWE Velociy model
    !! Do Cart -> RWE velocity Interpolation
    Iadj=.false.
    call C2R_init(0,0,norm,dx,dz,nx,nz,nt,nr,ox,oz)
    call C2R_run(RWEvel,raysin,CARTvel,gnorm,Iadj)  
    call sep_write(RWEvel,"velocity.H")

    !! Velocity
    do it=1,nt
       veltmp(it,1,1)=RWEvel(1,it)
    end do

    do ir=1,nr
       do it=2,nt  
          veltmp(it,1,ir) = (RWEvel(ir-1,it)+RWEvel(ir,it))/2.
       end do
       vmin(ir) = minval( RWEvel(ir,:) / flds(:,1,ir) )
    end do

    do iw=1,nw  !! Frequency Loop           

       Pimg=0.; sax=0.;

       w=ow+(iw-1)*dw
       do it=1,nt
          sax(it,1) = swf(it,iw)
       end do
       call rwetaper(sax) 

       do it=1,nt
          Pimg(it,1) = sax(it,1)
       end do
       !! Depth loop
       do ir=2,nr

          call shot_ker_onestep_noid(sax,veltmp(:,:,ir),flds(:,:,ir),w,-1.)
          call rwetaper(sax) 

          call fth(.false.,.false.,sax(:,1))
	  call wemig_phs_correction_onewf(sax(:,1),w,vmin(ir),-1.)
          call wemig_amp_correction_onewf(sax(:,1),w,veltmp(:,1,ir-1),veltmp(:,1,ir))
          call fth( .true.,.false.,sax(:,1) )
 
          call rwetaper(sax) 

          do it=1,nt
             Pimg(it,ir) = sax(it,1)
          end do

       end do !! . . z-loop
       write(0,*) w/(2.*3.1415),iw,nw,maxval(abs(Pimg))
       
       call sep_write(Pimg,"Rimage")

       !! Position wavefield on surface
       do ir = 1,nr
          if (near_rec(ir) .ge. nr) ix = nr-1
          if (near_rec(ir) .le. 0 ) ix = 1 
 
          ix = near_rec(ir)
          x = (ir-1)*rrd+rro
          dx1 = abs( x - raysin(ix,nt/2,1) )
          dx2 = abs(raysin(ix+1,nt/2,1) - x)
          dd1 = dx2/(dx1+dx2)
          dd2 = dx1/(dx1+dx2)
          swfout(ir,iw) = dd1*Pimg(nt/2,ix)+dd2*Pimg(nt/2,ix+1)
       end do
       
    end do  !! . . w-loop

    call shot_ker_release()
    call C2R_close()

  end subroutine wemig


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

  subroutine rwetaper_init()
    integer :: itap,jtap,ix
    integer :: 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(nt-itap+1) = cos(pi/2.* jtap/(ntap-1))
       end do
    end if
  end subroutine rwetaper_init

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

  subroutine rwetaper(dax)
    complex   :: dax(:,:)
    integer   :: it
    do it=1,nt
       dax(it,1) = dax(it,1) * tap(it)
    end do
  end subroutine rwetaper

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

  subroutine readin(  ax,az,aw,as,at,ar,rr)
    type(myaxis)   :: ax,az,aw,as,at,ar,rr

    !! Read in input dimensions
    call from_aux("swf","n1",at%n)
    call from_aux("swf","n2",aw%n)
    call from_aux("swf","n3",as%n)
    call from_aux("swf","d1",at%d)
    call from_aux("swf","d2",aw%d)
    call from_aux("swf","d3",as%d)
    call from_aux("swf","o1",at%o)
    call from_aux("swf","o2",aw%o)
    call from_aux("swf","o3",as%o)

    !! Read in input dimensions
    call from_aux("rwf","n1",rr%n)
    call from_aux("rwf","d1",rr%d)
    call from_aux("rwf","o1",rr%o)

    !! Read in input dimensions
    call from_aux("vel","n1",az%n)
    call from_aux("vel","n2",ax%n)
    call from_aux("vel","d1",az%d)
    call from_aux("vel","d2",ax%d)
    call from_aux("vel","o1",az%o)
    call from_aux("vel","o2",ax%o)

  end subroutine readin

  subroutine writeout( ax,az,aw,as,at,ar,rr)
    type(myaxis)    :: ax,az,aw,as,at,ar,rr

    !! . . Create offset axis
    call auxputch("n1","i",az%n,"image")
    call auxputch("n2","i",az%n,"image")
    call auxputch("n3","i",aw%n,"image")
    call auxputch("n4","i",as%n,"image")
    call auxputch("d1","r",az%d,"image")
    call auxputch("d2","r",ax%d,"image")
    call auxputch("d3","r",aw%d,"image")
    call auxputch("d4","r",as%d,"image")
    call auxputch("o1","r",az%o,"image")
    call auxputch("o2","r",ax%o,"image")
    call auxputch("o3","r",aw%o,"image")
    call auxputch("o4","r",as%o,"image")
    call auxputch("esize","i",8,"image")

    !! . . Output Imaged WF in RWE coordinates
    call auxputch("n1","i",at%n,"Rimage")
    call auxputch("n2","i",ar%n,"Rimage")
    call auxputch("n3","i",aw%n,"Rimage")
    call auxputch("n4","i",as%n,"Rimage")
    call auxputch("d1","r",at%d,"Rimage")
    call auxputch("d2","r",ar%d,"Rimage")
    call auxputch("d3","r",aw%d,"Rimage")
    call auxputch("d4","r",as%d,"Rimage")
    call auxputch("o1","r",at%o,"Rimage")
    call auxputch("o2","r",ar%o,"Rimage")
    call auxputch("o3","r",aw%o,"Rimage")
    call auxputch("o4","r",as%o,"Rimage")
    call auxputch("esize","i",8,"Rimage")

    !! . . Output Rays
    call auxputch("n1","i",ar%n,"rays.H")
    call auxputch("n2","i",at%n,"rays.H")
    call auxputch("d1","r",ar%d,"rays.H")
    call auxputch("d2","r",at%d,"rays.H")
    call auxputch("o1","r",ar%o,"rays.H")
    call auxputch("o2","r",at%o,"rays.H")
    call auxputch("esize","i",8,"rays.H")
    !! . . Output Rays
    call auxputch("n1","i",ar%n,"velocity.H")
    call auxputch("n2","i",at%n,"velocity.H")
    call auxputch("n3","i",2   ,"velocity.H")
    call auxputch("d1","r",ar%d,"velocity.H")
    call auxputch("d2","r",at%d,"velocity.H")
    call auxputch("o1","r",ar%o,"velocity.H")
    call auxputch("o2","r",at%o,"velocity.H")
    call auxputch("esize","i",4,"velocity.H")

    !! . . Output Surface Data
    call auxputch("n1","i",rr%n,"data")
    call auxputch("n2","i",aw%n,"data")
    call auxputch("n3","i",as%n,"data")
    call auxputch("d1","r",rr%d,"data")
    call auxputch("d2","r",aw%d,"data")
    call auxputch("d3","r",as%d,"data")
    call auxputch("o1","r",rr%o,"data")
    call auxputch("o2","r",aw%o,"data")
    call auxputch("o3","r",as%o,"data")
    call auxputch("esize","i",8,"data")

  end subroutine writeout
  
end module  RWE2D_source_wemig
