module shot_kernal_mod

!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 :: rax(:),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(:,:)
! 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=3.1415926
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 ( rax(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(rax,rbx,rcx,cax,cbx,ux,vx)
deallocate(ray,rby,rcy,cay,cby,uy,vy)
deallocate( xav)
end subroutine shot_ker_release

!input:vel: velocity(:,:)/a(:,:)
subroutine shot_ker_onestep(wld,vel,x,w,signn)
complex :: wld(:,:)
real :: vel(:,:),x(:,:),w
integer ::signn

integer ::ix,iy 

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

call shot_ker_ssf(wld,w,vel,x,ker_par,signn)
call shot_ker_fd(wld,xav/w,vel/w,ker_par,ker_par%ax1,ker_par%bx1,ker_par%ay1,ker_par%by1,signn)
call shot_ker_fd(wld,xav/w,vel/w,ker_par,ker_par%ax2,ker_par%bx2,ker_par%ay2,ker_par%by2,signn)

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
integer :: 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


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

integer :: nx,ny


integer :: ix,iy

nx=ker_par%nx; ny=ker_par%ny; tx=ker_par%tx; ty=ker_par%ty
!do ix=1,nx
!  uy=wld_z(ix,:)
!  cay=ca(ix,:)
!  cby=cb(ix,:)
!  call fd(uy,vy,ty,ay,by,cay,cby,ny,ray,rby,rcy)
!  wld_z(ix,:)=vy
!end do
!write(*,*) "maxval",maxval(real(wld_z))
do iy=1,ny
  ux=wld(:,iy)
  cax=ca(:,iy)
  cbx=cb(:,iy)
  call fd(ux,vx,tx,ax,bx,cax,cbx,nx,rax,rbx,rcx)
  wld(:,iy)=vx
end do
end subroutine shot_ker_fd

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

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



end module shot_kernal_mod
