module Kernel_gogi
  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

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

  complex,allocatable :: contablepp(:,:,:,:),contablepm(:,:,:,:),conapp(:),conapm(:)

  ! for mult-ref-vel
  real, allocatable :: v0_z(:)
  complex,allocatable :: tmp_wld(:,:)
  real,allocatable :: xav(:,:),vold(:),vnew(:),Acor(:)
  ! end for mult-ref-vel

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))

    !! For amp filter
    allocate( vold(nx),vnew(nx),Acor(nx))

  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(id,wld,vel,x,w,signn)
    complex :: wld(:,:,:)
    real :: vel(:,:),x(:,:),w
    real ::signn

    integer ::ix,iy,id

    do iy=1,ker_par%ny
       do ix=1,ker_par%nx
          xav(ix,iy)=vel(ix,iy)/x(ix,iy)
       end do
    end do

    call shot_ker_ssf(id,wld,w,vel,x,ker_par,signn)

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

    !! . . Second FFD step
    call shot_ker_fd (id,wld,xav/w,xav/w,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_onestep_noid(wld,vel,x,w,signn)
    complex :: wld(:,:)
    real :: vel(:,:),x(:,:),w
    real ::signn

    integer ::ix,iy

    do iy=1,ker_par%ny
       do ix=1,ker_par%nx
          xav(ix,iy)=vel(ix,iy)/x(ix,iy)
       end do
    end do

    call shot_ker_ssf_noid(wld,w,vel,x,ker_par,signn)

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

    !! . . Second FFD step
    call shot_ker_fd_noid (wld,xav/w,xav/w,ker_par,-signn*ker_par%ax2,ker_par%bx2,ker_par%ay2,ker_par%by2)
   xav=0.
  end subroutine shot_ker_onestep_noid
!###############################################
  subroutine shot_ker_ssf(id,wld,w,vel,x,ker_par,signn)
    complex :: wld(:,:,:)
    real    :: w,vel(:,:),x(:,:)
    type(ker_par_type) :: ker_par
    real :: signn

    integer :: nx,ny
    real    :: dz

    integer :: ix,iy,id
    real    :: phsft

    nx=ker_par%nx; ny=ker_par%ny; dz=ker_par%dz
    do iy=1,ny
       do ix=1,nx
          phsft=-signn*dz*(x(ix,iy)*w/vel(ix,iy))
          wld(ix,iy,id)=wld(ix,iy,id)*cmplx(cos(phsft),sin(phsft))
       end do
    end do

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

    nx=ker_par%nx; ny=ker_par%ny; dz=ker_par%dz
    do iy=1,ny
       do ix=1,nx
          phsft=-signn*dz*(x(ix,iy)*w/vel(ix,iy))
          wld(ix,iy)=wld(ix,iy)*cmplx(cos(phsft),sin(phsft))
       end do
    end do

  end subroutine shot_ker_ssf_noid

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

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

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

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

    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_noid

!#############################################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,signn
    integer :: nx

    integer :: ix
    complex :: url,urr,ror,rol

    do ix=1,nx
       if (ix==1)  then
          url=0.0
       else
          url=u(ix-1)
       end if
       if (ix==nx) then
          urr=0.0
       else
          urr=u(ix+1)
       end if
       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
    call thryj(ra,rb,rc,v,nx,1)
  end subroutine fd

  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(id,sdax,rdax,w,vmin,sarg,rarg)
    integer :: ir,inn,it,id
    real    :: w,kr,sarg,rarg,carg,carg1,vmin,pkr
    complex :: sdax(:,:),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 
          sdax(ir,id) = 0.
          rdax(ir,id) = 0. 
       else
          phsft=-dz*w_v0*(sqrt(1-sr2)-(1+ (a1*sr2)/(1+b1*sr2)+(a2*sr2)/(1+b2*sr2) ) )
          sdax(ir,id)=sdax(ir,id)*cmplx(cos(phsft),sin(phsft))
          rdax(ir,id)=rdax(ir,id)*cmplx(cos(phsft),-sin(phsft)) 
       end if
    end do  

  end subroutine wemig_phs_correction

  subroutine wemig_phs_correction_onewf(sdax,w,vmin,sarg)
    integer :: ir,inn,it
    real    :: w,kr,sarg,carg,carg1,vmin,pkr
    complex :: sdax(:)
    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 
          sdax(ir) = 0.
       else
          phsft=-dz*w_v0*(sqrt(1-sr2)-(1+ (a1*sr2)/(1+b1*sr2)+(a2*sr2)/(1+b2*sr2) ) )
          sdax(ir)=sdax(ir)*cmplx(cos(phsft),sin(phsft))
       end if
    end do  

  end subroutine wemig_phs_correction_onewf

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!
!! . . Amplitude correction
!!
  subroutine wemig_amp_correction_onewf(sdax,w,vold,vnew)
    integer :: ir,inn,it
    real    :: w,kt,pkr,vold(:),vnew(:)!,Acor(:)
    complex :: sdax(:)
    real    :: w_v0,sr,sr2,phsft,dz,a1,a2,b1,b2

    dz=ker_par%dz

    do it=1,ker_par%nx
       kt=ker_par%fkx+(ir-1)*ker_par%dkx  

       Acor(it) = - ( sqrt(1./vnew(it)**2 - w*kt) - sqrt(1./vold(it)**2 - w*kt) ) / &
              (2.* dz * sqrt( 1./ (vnew(it)/2.+vold(it)/2.)**2-w*kt) ) 

       sdax(it)=sdax(it) * ( 1. + dz*Acor(it) )

    end do  

  end subroutine wemig_amp_correction_onewf


  !----------------------------------------------------------------  
  !!
  !! Fourier-domain phase shift
  !!
  subroutine correction(id,rdax,w,vmin,rarg)
    integer :: ir,inn,it,id
    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

    !carg1= (w/vmin)**2
    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;
       !carg = carg1 - kr**2

       if ( sr2>1) then 
          rdax(ir,id) = 0. 
       else
          phsft=-dz*w_v0*(sqrt(1-sr2)-(1+ (a1*sr2)/(1+b1*sr2)+(a2*sr2)/(1+b2*sr2) ) )
          rdax(ir,id)=rdax(ir,id)*cmplx(cos(phsft),-sin(phsft)) 
       end if
    end do  

  end subroutine correction
 
end module Kernel_gogi
