!!---------------------------------------------------------------------------
!! Author: Yaxun Tang, tang@sep.stanford.edu
!!---------------------------------------------------------------------------
module wem_ssrepi_mod
  implicit none

contains
  !----------------------------------------------------------------------------------------------------------------------
  !.. Multiple reference velocity extrapolation plus interpolation
  !----------------------------------------------------------------------------------------------------------------------
  integer function ssr_epi(Fxop, wfld_slice, vmod_slice, vref_slice, kx2, ky2, ws, dz, nx, ny, nvrf, opsign, &
                           iwrk2d, cwrk3d, plan_2dxF, plan_2dxI) result(stat)

    complex, dimension(:,:)   :: wfld_slice
    real   , dimension(:,:)   :: vmod_slice
    real   , dimension(:)     :: vref_slice
    real   , dimension(:)     :: kx2, ky2
    real                      :: ws
    real                      :: dz
    integer                   :: nx, ny, nvrf
    integer                   :: opsign
    integer, dimension(:,:)   :: iwrk2d
    complex, dimension(:,:,:) :: cwrk3d
    integer*8                 :: plan_2dxF
    integer*8                 :: plan_2dxI

    interface
      integer function Fxop(wfld_slice, vmod_slice, kx2, ky2, vref, ws, dz, nx, ny, opsign, plan_2dxF, plan_2dxI) result(stat)
        complex, dimension(:,:) :: wfld_slice
        real   , dimension(:,:) :: vmod_slice
        real   , dimension(:)   :: kx2, ky2
        real                    :: vref
        real                    :: ws
        real                    :: dz
        integer                 :: nx, ny
        integer                 :: opsign
        integer*8               :: plan_2dxF
        integer*8               :: plan_2dxI
      end function
    end interface

    integer                     :: ix, iy, itr, ivrf
    real                        :: vmin, vmax, dvel, vref

    if (nvrf == 1) then
       vref = vref_slice(1)
       stat = FXop(wfld_slice, vmod_slice, kx2, ky2, vref, ws, dz, nx, ny, opsign, plan_2dxF, plan_2dxI)

    else

       iwrk2d = 0
       cwrk3d = 0. 

       vmin = minval(vmod_slice)
       vmax = maxval(vmod_slice)
       dvel = (vmax - vmin) / (nvrf - 1)

       do iy=1, ny
       do ix=1, nx
          iwrk2d(ix,iy) = 1.5 + (vmod_slice(ix, iy)-vmin)/dvel
          iwrk2d(ix,iy) = max(iwrk2d(ix,iy),    1)
          iwrk2d(ix,iy) = min(iwrk2d(ix,iy), nvrf)
       end do
       end do

       cwrk3d(:,:,1) = wfld_slice
       cwrk3d(:,:,2) = 0.
       do ivrf=1, nvrf
          wfld_slice = cwrk3d(:,:,1)
          vref = vref_slice(ivrf)
          stat = FXop(wfld_slice, vmod_slice, kx2, ky2, vref, ws, dz, nx, ny, opsign, plan_2dxF, plan_2dxI)
          where (iwrk2d == ivrf)
                cwrk3d(:,:,2) = wfld_slice
          end where
       end do

       wfld_slice = cwrk3d(:,:,2)

   end if
   stat = 0
  end function


end module
