!!.. Differential residual migration
!!.. Kw(rho) = Kw(rho0)+(rho-rho0)dKw/drho|rho=rho0
!!.. This is the Fourier domain operation part
!!.. the other part should be applied in the space domain with -iz
!!.. for simplicity, here I linearize it around rho0=1
module mva_resdif_mod
  use  base_struct_mod
  use  base_cstret_mod
  implicit none

  real   , dimension(:), pointer    , private :: kx , ky
  real                              , private :: okz, dkz
  integer                           , private :: nkx, nky, nkz
  complex, dimension(:), allocatable, private :: tri, tro
  real   , dimension(:), allocatable, private :: str


contains


  integer function mva_resdif_init(kx_in, ky_in, akz, eps) result(stat)
    real, dimension(:), target :: kx_in, ky_in
    type(axis)                 :: akz
    real                       :: eps

    stat = 1
    kx => kx_in
    ky => ky_in

    nkx = size(kx); nky = size(ky)
    nkz = akz%n; okz = akz%o; dkz = akz%d

    stat = base_cstret_init(nkz, nkz, okz, dkz, eps)

    allocate(tri(nkz), tro(nkz), str(nkz))

    stat = 0
  end function

  integer function mva_resdif_clos() result(stat)
    stat = 1
    nullify(kx)
    nullify(ky)
    deallocate(tri, tro, str)
    stat = base_cstret_clos()
    stat = 0
  end function

  integer function mva_resdif_oper(imgi, imgo, rho, khx, khy) result(stat)
    complex, dimension(:,:,:) :: imgi
    complex, dimension(:,:,:) :: imgo
    real                      :: rho
    real                      :: khx, khy

    integer                   :: ikx, iky, ikz
    real                      :: mks, mkg, kz, sca
    real                      :: kzh, kzm, zzs, zzg

write(0,*)"rho, khx, khy", rho, khx, khy

    imgo = 0.

    !!.. differential residual migration
    do iky=1, nky
    do ikx=1, nkx
       mks = sqrt((kx(ikx)-khx)**2+(ky(iky)-khy)**2)  !!.. |ks|
       mkg = sqrt((kx(ikx)+khx)**2+(ky(iky)+khy)**2)  !!.. |kg|
       do ikz=2, nkz 
          kz = okz + (ikz-1)*dkz
          kzh = (4*kz**2+(mkg-mks)**2)/4.
          kzm = (4*kz**2+(mkg+mks)**2)/4.
          zzs = 1./rho**2 * kzh*kzm - kz**2 * mks**2
          zzg = 1./rho**2 * kzh*kzm - kz**2 * mkg**2

          if (zzs>0. .and. zzg>0.) then
             sca = 0.5/kz*(1./sqrt(zzs)+1./sqrt(zzg))*kzh*kzm/rho
          else
             sca = -2*dkz
          end if

          imgo(ikz,ikx,iky) = imgi(ikz,ikx,iky)*sca

       end do
    end do
    end do

!write(0,*)maxval(cabs(imgo)), minval(cabs(imgo))


    imgi = imgo

   !!.. residual migration
    do iky=1, nky
    do ikx=1, nkx
       mks = sqrt((kx(ikx)-khx)**2+(ky(iky)-khy)**2)  !!.. |ks|
       mkg = sqrt((kx(ikx)+khx)**2+(ky(iky)+khy)**2)  !!.. |kg|
       do ikz=2, nkz
          kz = okz + (ikz-1)*dkz
          kzh = (4.*kz**2+(mkg-mks)**2)/4.
          kzm = (4.*kz**2+(mkg+mks)**2)/4.
          zzs = 1./rho**2 * kzh*kzm - kz**2 * mks**2
          zzg = 1./rho**2 * kzh*kzm - kz**2 * mkg**2
          if (zzs>0. .and. zzg>0) then
             str(ikz) = 0.5/kz*(sqrt(zzs)+sqrt(zzg))
          else
             str(ikz) = -2*dkz
          end if
       end do
       tri = imgi(:,ikx,iky)
       stat = base_cstret_oper(str, tri, tro)
       imgo(:,ikx,iky) = tro
    end do
    end do

    stat = 0
  end function


end module 
