module Kernel_TOPO
  use sep

  !parameters for shot_phase_shift
  implicit none
  type ker_par_type
     integer :: nx,ny
     real    :: dx,dy,dz
     real    :: fkx,dkx,fky,dky   ! for phase-shift
     real    :: ax1,ay1,bx1,by1,ax2,bx2,ay2,by2,tx,ty ! for finite difference
  end type ker_par_type

  type(ker_par_type),private:: ker_par

  complex,allocatable :: dax(:),rbx(:),rcx(:),ux(:),vx(:),ray(:),rby(:),rcy(:),uy(:),vy(:)
  real,allocatable    :: cax(:),cbx(:), cay(:),cby(:)

  real, allocatable :: v0_z(:),xav(:,:),vold(:),vnew(:)
  complex,allocatable :: tmp_wld(:,:)

contains

  !###############################################################

  subroutine shot_ker_init(nx,ny,dx,dy,dz,a,b,trick)
    integer :: nx,ny
    real    :: dx,dy,dz,a,b,trick
    real    :: pi
    real    :: a1,a2,b1,b2
    a1=0.040315157; b1=0.873981642
    a2=0.457289566; b2=0.222691983
    pi=acos(-1.)
    ker_par%nx=nx; ker_par%ny=ny
    ker_par%dx=abs(dx); ker_par%dy=abs(dy); ker_par%dz=dz
    ker_par%fkx=-pi/ker_par%dx; ker_par%fky=-pi/ker_par%dy
    ker_par%dkx=2.0*pi/(nx*ker_par%dx); ker_par%dky=2.0*pi/(ny*ker_par%dy)

    ker_par%ax1=a1*dz/(2.0*ker_par%dx*ker_par%dx); 
    ker_par%ay1=a1*dz/(2.0*ker_par%dy*ker_par%dy)
    ker_par%bx1=b1/(ker_par%dx*ker_par%dx); 
    ker_par%by1=b1/(ker_par%dy*ker_par%dy)

    ker_par%ax2=a2*dz/(2.0*ker_par%dx*ker_par%dx); 
    ker_par%ay2=a2*dz/(2.0*ker_par%dy*ker_par%dy)
    ker_par%bx2=b2/(ker_par%dx*ker_par%dx); 
    ker_par%by2=b2/(ker_par%dy*ker_par%dy)

    ker_par%tx=trick;     ker_par%ty=trick
    allocate ( dax(nx),rbx(nx),rcx(nx),ux(nx),vx(nx),cax(nx),cbx(nx) )
    allocate ( ray(ny),rby(ny),rcy(ny),uy(ny),vy(ny),cay(ny),cby(ny) )
    allocate ( xav(nx,ny) )

  end subroutine shot_ker_init

  !###############################################################

  subroutine shot_ker_release
    deallocate(dax,rbx,rcx,cax,cbx,ux,vx)
    deallocate(ray,rby,rcy,cay,cby,uy,vy)
    deallocate( xav)
  end subroutine shot_ker_release

  !###############################################################

  subroutine shot_ker_onestep(wld,vel,x,w,signn)
    integer ::ix,iy
    real :: vel(:,:),x(:,:),w,signn
    complex :: wld(:,:)

    do iy=1,ker_par%ny
       do ix=1,ker_par%nx
          xav(ix,iy)= sqrt( abs((x(ix,2)*x(ix,3))**2 / &
                 ( (w * x(ix,2) / vel(ix,iy))**2 - (x(ix,4)/w)**2 )) )
       end do
    end do

    !! . . Initial space shift
    call shot_ker_ssf(wld,w,vel,x,ker_par,signn)

    !! . . First FFD step
    call shot_ker_fd (wld,xav,xav,ker_par,-signn*ker_par%ax1,ker_par%bx1,ker_par%ay1,ker_par%by1)

    !! . . Second FFD step
    call shot_ker_fd (wld,xav,xav,ker_par,-signn*ker_par%ax2,ker_par%bx2,ker_par%ay2,ker_par%by2)
    xav=0.
  end subroutine shot_ker_onestep


  !###############################################
  subroutine shot_ker_ssf(wld,w,vel,x,ker_par,signn)
    complex :: wld(:,:)
    real    :: w,vel(:,:),x(:,:)
    type(ker_par_type) :: ker_par
    real :: signn
    integer :: nx,ny,ix,iy
    real    :: dz,phsft

    nx=ker_par%nx; ny=ker_par%ny; dz=ker_par%dz

    do iy=1,ny
       do ix=1,nx

          phsft=-signn*dz * sqrt( abs(( x(ix,2) * w / vel(ix,iy))**2 - x(ix,4)**2) )
          wld(ix,iy)=exp(signn*dz*x(ix,1))*wld(ix,iy)*&
          cmplx(cos(phsft),sin(phsft))

       end do
    end do

  end subroutine shot_ker_ssf


  !#######################################

  subroutine shot_ker_fd(wld,ca,cb,ker_par,ax,bx,ay,by)
    integer :: nx,ny,ix,iy
    real    :: ca(:,:),cb(:,:),tx,ty,ax,ay,bx,by
    complex :: wld(:,:)
    type(ker_par_type) ::ker_par 

    nx=ker_par%nx; ny=ker_par%ny; tx=ker_par%tx; ty=ker_par%ty
    do iy=1,ny
       ux=wld(:,iy)
       cax=ca(:,iy)
       cbx=cb(:,iy)
       call fd(ux,vx,tx,ax,bx,cax,cbx,nx,dax,rbx,rcx)
       wld(:,iy)=vx
    end do
  end subroutine shot_ker_fd

  !#############################################3

  subroutine fd(u,v,tr,a,b,ca,cb,nx,ra,rb,rc)
    complex :: u(:),v(:)
    complex :: ra(:),rb(:),rc(:)
    real    :: ca(:),cb(:)
    real    :: a,b,tr
    integer :: nx

    integer :: ix
    complex :: url,urr,ror,rol
    !! Case where ix=1
    url=0.
    urr=u(2)
    ror=cmplx(tr+b*cb(1)*cb(1),a*ca(1))
    rol=conjg(ror)
    v(1)=u(1)+ror*(url+urr-2.0*u(1))
    ra(1)=1.0-2.0*rol
    rb(1)=rol
    rc(1)=0.

    do ix=2,nx-1
        url=u(ix-1)
       urr=u(ix+1)
       ror=cmplx(tr+b*cb(ix)*cb(ix),a*ca(ix))
       rol=conjg(ror)
       v(ix)=u(ix)+ror*(url+urr-2.0*u(ix))
       ra(ix)=1.0-2.0*rol
       if (ix <nx) rb(ix) =rol
       if (ix >1) rc(ix-1)=rol
    end do

    !! Case where ix=nx
    url=u(nx-1)
    urr=0.
    ror=cmplx(tr+b*cb(nx)*cb(nx),a*ca(nx))
    rol=conjg(ror)
    v(nx)=u(nx)+ror*(url+urr-2.0*u(nx))
    ra(nx)=1.0-2.0*rol
    rb(nx)=0.
    rc(nx-1)=rol
   call thryj(ra,rb,rc,v,nx,1)

  end subroutine fd

  !#############################################3

  subroutine thryj(a,b,c,v,h,yJ)
    implicit none
    integer h,yJ
    complex :: a(:),b(:),c(:),v(:)
    integer :: i
    do i=1,yJ
       b(i)=b(i)/a(i)
       v(i)=v(i)/a(i)
    enddo
    do i=yJ+1,h
       a(i)=a(i)-b(i-yJ)*c(i-yJ)
       v(i)=(v(i)-c(i-yJ)*v(i-yJ))/a(i)
       if(i.le.h-yJ) then
          b(i)=b(i)/a(i)
       endif
    enddo
    do i=h-yJ,1,-1
       v(i)=v(i)-b(i)*v(i+yJ)
    enddo
  end subroutine thryj

  !----------------------------------------------------------------  
  !!
  !! Fourier-domain phase shift
  !!
  subroutine wemig_phs_correction(rdax,w,vmin,rarg)
    integer :: ir,inn,it
    real    :: w,kr,rarg,carg,carg1,vmin,pkr
    complex :: rdax(:,:)
    real    :: w_v0,sr,sr2,phsft,dz,a1,a2,b1,b2
    a1=-0.040315157; b1=-0.873981642
    a2=-0.457289566; b2=-0.222691983
    dz=ker_par%dz
    w_v0=w/vmin

    do ir=1,ker_par%nx
       kr=ker_par%fkx+(ir-1)*ker_par%dkx  
       sr=kr/w_v0;
       sr2=sr*sr;

       if ( sr2>1) then 
          rdax(ir,1) = 0. 
          rdax(ker_par%nx-ir+1,1)=0.
       else
          phsft=-dz*w_v0*(sqrt(1-sr2)-(1+(a1*sr2)/(1+b1*sr2)+(a2*sr2)/(1+b2*sr2)))

          rdax(ir,1)=rdax(ir,1)*cmplx(cos(phsft),-sin(phsft)) 

          rdax(ker_par%nx-ir+1,1)=&
                      rdax(ker_par%nx-ir+1,1)*cmplx(cos(phsft),-sin(phsft)) 
       end if
    end do

  end subroutine wemig_phs_correction


end module Kernel_TOPO
