!! wavefiled extrapolation with multiple reference velocities
!! extrapolate for one frequency and one depth step
!! should initialize for each frequency and depth step
module wei_multi_ref_kernel_mod
  implicit none

  integer, private                            :: nref 
  integer, dimension(:,:), pointer, private   :: index_vref
  complex, dimension(:,:), pointer, private   :: wfld_slice_tmp1
  complex, dimension(:,:), pointer, private   :: wfld_slice_tmp2

contains

!============================================================================================
  integer function wei_multi_ref_kernel_init(kx2_in, ky2_in, dz_in, maxang_in, nref_in, FXin) result(stat)
    real, dimension(:), pointer     :: kx2_in
    real, dimension(:), pointer     :: ky2_in
    real                            :: dz_in
    real                            :: maxang_in
    integer                         :: nref_in
    
    interface
      integer function FXin(kx2_in, ky2_in, dz_in, maxang_in) result(stat)
        real, dimension(:), pointer :: kx2_in
        real, dimension(:), pointer :: ky2_in
        real                        :: dz_in
        real                        :: maxang_in
      end function
    end interface

    integer                         :: err
    stat = 1

    nref = nref_in
    stat = FXin(kx2_in, ky2_in, dz_in, maxang_in)
    allocate(index_vref(size(kx2_in), size(ky2_in)))
    allocate(wfld_slice_tmp1(size(kx2_in), size(ky2_in)))
    allocate(wfld_slice_tmp2(size(kx2_in), size(ky2_in)))

    stat = 0
  end function
!=============================================================================================
  integer function wei_multi_ref_kernel_oper(adj, opsign, ws, vref_slice, vmod_slice, wfld_slice, FXop) result(stat)
    logical                          :: adj
    integer                          :: opsign
    real                             :: ws
    real   , dimension(:)  , pointer :: vref_slice
    real   , dimension(:,:), pointer :: vmod_slice
    complex, dimension(:,:), pointer :: wfld_slice

    interface
      integer function FXop(adj, opsign, ws, vref, vmod_slice, wfld_slice, vzref) result(stat)
        optional                         :: vzref
        logical                          :: adj
        integer                          :: opsign
        real                             :: ws
        real                             :: vref
        real                             :: vzref
        real   , dimension(:,:), pointer :: vmod_slice
        complex, dimension(:,:), pointer :: wfld_slice
      end function
    end interface


    integer                          :: iref, ix, iy
    real                             :: vref, vmax, vmin, dvel

    if (nref == 1) then
         vref = vref_slice(1)
         stat = FXop(adj, opsign, ws, vref, vmod_slice, wfld_slice)
    else
        index_vref = 0
        wfld_slice_tmp1 = 0.
        wfld_slice_tmp2 = 0.
        vmin = minval(vmod_slice)
        vmax = maxval(vmod_slice)
        dvel = (vmax - vmin) / (nref - 1)
        do iy=1, size(vmod_slice, 2)
           do ix=1, size(vmod_slice, 1)
              index_vref(ix, iy) = max(1, int(1.5 + (vmod_slice(ix, iy) - vmin)/dvel))
              index_vref(ix, iy) = min(index_vref(ix, iy), nref)
           end do
        end do

        wfld_slice_tmp1 = wfld_slice
        wfld_slice = 0.
        do iref=1, nref
           wfld_slice_tmp2 = wfld_slice_tmp1
           vref = vref_slice(iref)
           stat = FXop(adj, opsign, ws, vref, vmod_slice, wfld_slice_tmp2)
           where (index_vref == iref)
                 wfld_slice = wfld_slice + wfld_slice_tmp2
           end where
        end do
    end if
    
  end function
!=============================================================================================
  integer function wei_multi_ref_kernel_clos(FXnd) result(stat)
    interface
      integer function FXnd() result(stat)
      end function
    end interface

    stat = FXnd()
    if(associated(index_vref))      deallocate(index_vref)
    if(associated(wfld_slice_tmp2)) deallocate(wfld_slice_tmp1)
    if(associated(wfld_slice_tmp2)) deallocate(wfld_slice_tmp2)
    stat = 0
  end function

end module







