!!.. version 3.0
!!.. This version won't output source and receiver wavefields and 
!!.. is faster since I/O time has been saved
!!.. This version improves the imaging condition
!!.. This version uses complex image
module wemva_kerpix_complex_mod
  use  base_blkcal_mod
  use  base_ioutil_mod
  use  base_spline_mod
  use  wem_velref_mod
  use  wem_wfdexp_mod
  use  wem_imgcon_complex_mod
  use  wem_tapers_mod
  use  wem_fftwv3_mod
  use  wemva_struct_mod
  use  wemva_histos_complex_mod
  use  wemva_glocal_mod
  use  wemva_fileio_complex_mod
  use  wemva_bornsc_mod
  implicit none

  type(dat_grid)                              , private :: dgrid    !!.. receiver grid
  type(vel_grid)                              , private :: vgrid    !!.. velocity grid
  type(pro_grid)                              , private :: pgrid    !!.. propagation grid
  type(img_grid)                              , private :: igrid    !!.. image grid
  type(out_grid)                              , private :: ogrid    !!.. image output grid

  type(block)                                 , private :: wblock

  type(wemva_ftag)                            , private :: ftag_

  complex  , dimension(:,:)      , allocatable, private :: csou, crec
  complex  , dimension(:,:,:,:)  , allocatable, private :: bwfs, bwfr
  complex  , dimension(:,:,:,:)  , allocatable, private :: dwfs, dwfr
  complex  , dimension(:,:,:)    , allocatable, private :: bwfs_slice, bwfr_slice
  complex  , dimension(:,:,:)    , allocatable, private :: dwfs_slice, dwfr_slice
  complex  , dimension(:,:,:)    , allocatable, private :: wtmp_slice

  real     , dimension(:,:)      , allocatable, private :: bvel
  complex  , dimension(:,:)      , allocatable, private :: cvel 
  complex  , dimension(:,:,:,:,:), allocatable, private :: bimg
  complex  , dimension(:,:,:)    , allocatable, private :: dslw
  complex  , dimension(:,:,:,:,:), allocatable, private :: dimg
  real     , dimension(:)        , allocatable, private :: bvrf

  real     , dimension(:)        , allocatable, private :: kx2, ky2, tapx, tapy

  integer  , dimension(:,:,:)    , allocatable, private :: iwrk2d
  real     , dimension(:,:)      , allocatable, private :: rwrk1d
  complex  , dimension(:,:)      , allocatable, private :: cwrk1d
  complex  , dimension(:,:,:,:)  , allocatable, private :: cwrk3d

  integer*8, dimension(:)        , allocatable, private :: plan_2dsF, plan_2dsI
  integer*8, dimension(:)        , allocatable, private :: plan_2drF, plan_2drI
  integer*8, dimension(:)        , allocatable, private :: plan_2dtF, plan_2dtI

  integer  , dimension(:,:)      , allocatable, private :: index_mxs, index_mxr
  integer  , dimension(:,:)      , allocatable, private :: index_mys, index_myr

  integer                                     , private :: nopbsc, typbsc, nzrept
  real                                        , private :: etabsc

  integer                                     , private :: ompchunk, nth
  integer                           , external, private :: omp_get_num_procs
  integer                           , external, private :: omp_get_num_threads
  integer                           , external, private :: omp_set_num_threads
  integer                           , external, private :: omp_get_thread_num

  real     ,                         parameter, private :: PI = 3.14159265359
contains

  integer function wemva_kerpix_init(vgrid_out, igrid_out, ftag_out_) result(stat)
    type(vel_grid)   :: vgrid_out
    type(img_grid)   :: igrid_out
    type(wemva_ftag) :: ftag_out_

    ftag_%csou_ = "csou"
    ftag_%crec_ = "crec"
    ftag_%hsou_ = "hsou"
    ftag_%bslw_ = "bslw"
    ftag_%bimg_ = "bimg"
    ftag_%dslw_ = "dslw"
    ftag_%dimg_ = "dimg"

    call auxinout(ftag_%csou_)
    call auxinout(ftag_%crec_)
    call auxinout(ftag_%bslw_)
    call auxinout(ftag_%dslw_)
    call auxinout(ftag_%dimg_)

    call from_param("radian", dgrid%radian, .true. )
    call from_param("offset", dgrid%offset, .false.)
    call from_param("comsou", dgrid%comsou, .false.)
    call from_param("hdrsou", dgrid%hdrsou, .false.)
    call from_param("wntimg", igrid%wntimg, .false.)
    call from_param("ntpx"  , pgrid%ntpx  , 1      )
    call from_param("ntpy"  , pgrid%ntpy  , 1      )
    call from_param("nvrf"  , pgrid%nvrf  , 1      )
    call from_param("pdip"  , pgrid%pdip  , 65     )
    call from_param("oper"  , pgrid%oper  , "ssfpi")
    call from_param("trick" , pgrid%trick , 0.14   )
    call from_param("nws"   , wblock%ns   , 1      )
    call from_param("nopbsc", nopbsc      , 1      )
    call from_param("typbsc", typbsc      , 1      )
    call from_param("etabsc", etabsc      , 0.01   )
    call from_param("nzrept", nzrept      , 100    )

    call wemva_histo_data(dgrid, vgrid, igrid, pgrid, ftag_)

    ompchunk = 1
    nth      = 1

    !$OMP PARALLEL
    if (omp_get_thread_num() == 0) then
       nth = omp_get_num_threads()
       write(0,*)"nthread =", nth, omp_get_num_procs()
    end if
    !$OMP END PARALLEL

    !!.. if present the source header, read it
    if (dgrid%hdrsou) then
       allocate(dgrid%hsou(2,dgrid%nshg))
       call sep_read(dgrid%hsou, ftag_%hsou_)
       write(0,*)"source locations:", dgrid%hsou
    end if

    if (.not.dgrid%radian) then
       dgrid%aw%d = dgrid%aw%d*2*PI; dgrid%aw%o = dgrid%aw%o*2*PI
    end if

    call wemva_glocal(pgrid, ogrid, vgrid, dgrid, igrid, 1)
    call base_blkcal(wblock, dgrid%aw)

    write(0,*)"oper = ", pgrid%oper
    write(0,*)"wblock%nb", wblock%nb, wblock%blen

    allocate(csou(dgrid%ntrs,wblock%ns))
    allocate(crec(dgrid%ntrg,wblock%ns))
    allocate(bwfs(pgrid%ax%n,pgrid%ay%n,wblock%ns,pgrid%az%n))
    allocate(bwfr(pgrid%ax%n,pgrid%ay%n,wblock%ns,pgrid%az%n))
    allocate(dwfs(pgrid%ax%n,pgrid%ay%n,wblock%ns,pgrid%az%n))
    allocate(dwfr(pgrid%ax%n,pgrid%ay%n,wblock%ns,pgrid%az%n))
    allocate(cvel(pgrid%ax%n,pgrid%ay%n))
    allocate(bvel(pgrid%ax%n,pgrid%ay%n))
    allocate(bvrf(pgrid%nvrf))
    allocate(dslw(pgrid%ax%n,pgrid%ay%n,nth))
    allocate(dimg(ogrid%ahx%n,ogrid%ahy%n,ogrid%ax%n,ogrid%ay%n,nth))
    allocate(bwfs_slice(pgrid%ax%n,pgrid%ay%n,nth))
    allocate(bwfr_slice(pgrid%ax%n,pgrid%ay%n,nth))
    allocate(dwfs_slice(pgrid%ax%n,pgrid%ay%n,nth))
    allocate(dwfr_slice(pgrid%ax%n,pgrid%ay%n,nth))
    allocate(wtmp_slice(pgrid%ax%n,pgrid%ay%n,nth))
    allocate(kx2(pgrid%ax%n),ky2(pgrid%ay%n))
    allocate(tapx(pgrid%ntpx),tapy(pgrid%ntpy))

    allocate(plan_2dsF(nth),plan_2dsI(nth))
    allocate(plan_2drF(nth),plan_2drI(nth))
    allocate(plan_2dtF(nth),plan_2dtI(nth))

    allocate(index_mxs(ogrid%ahx%n,ogrid%ax%n))
    allocate(index_mxr(ogrid%ahx%n,ogrid%ax%n))
    allocate(index_mys(ogrid%ahy%n,ogrid%ay%n))
    allocate(index_myr(ogrid%ahy%n,ogrid%ay%n))

    allocate(cwrk1d(pgrid%ax%n*5+pgrid%ay%n*5,nth)) !!.. ifd or ffd
    allocate(cwrk3d(pgrid%ax%n,pgrid%ay%n,2,nth))   !!.. ssf-epi
    allocate(iwrk2d(pgrid%ax%n,pgrid%ay%n,nth))     !!.. ssf-epi
    allocate(rwrk1d(pgrid%ax%n+pgrid%ay%n,nth))

    call k2_map(kx2, pgrid%ax%n, pgrid%ax%d)
    call k2_map(ky2, pgrid%ay%n, pgrid%ay%d)

    call wem_fftw_init(plan_2dsF, plan_2dsI, dwfs_slice, pgrid%ax%n, pgrid%ay%n, nth)
    call wem_fftw_init(plan_2drF, plan_2drI, dwfr_slice, pgrid%ax%n, pgrid%ay%n, nth)
    call wem_fftw_init(plan_2dtF, plan_2dtI, wtmp_slice, pgrid%ax%n, pgrid%ay%n, nth)

    stat = wem_tapers_init(tapx, tapy, pgrid%ntpx, pgrid%ntpy)

    !!.. initilize backgorund image if requested
    if (igrid%wntimg) then
       call auxinout(ftag_%bimg_)
       allocate(bimg(ogrid%ahx%n,ogrid%ahy%n,ogrid%ax%n,ogrid%ay%n,nth))
       call wemva_histo_imag(igrid, ftag_%bimg_)
       call wemva_fio_imag_init(igrid, ftag_%bimg_)
       write(0,*)"generate background image"
    end if 

    vgrid_out = vgrid
    igrid_out = igrid
    ftag_out_ = ftag_

    stat = 0
  end function
  !!-----------------------------------------------------------------------------------------
  integer function wemva_kerhis_init(adj, dslw_, dimg_) result(stat)
    logical          :: adj
    character(len=*) :: dslw_, dimg_
    if (adj) then
       call wemva_histo_dslw(vgrid, dslw_)
    else
       call wemva_histo_imag(igrid, dimg_)
    end if
    stat = 0
  end function
  !!-----------------------------------------------------------------------------------------
  integer function wemva_kerpix_clos() result(stat)
    stat = 1
    call auxclose(ftag_%csou_)
    call auxclose(ftag_%crec_)
    call auxclose(ftag_%bslw_)
    call auxclose(ftag_%dslw_)
    call auxclose(ftag_%dimg_)
    if (igrid%wntimg) call auxclose(ftag_%bimg_)

    deallocate(csou, crec, cvel)
    deallocate(bwfs, bwfr, dwfs, dwfr)
    deallocate(bvel, bvrf, dslw, dimg)
    deallocate(bwfs_slice, bwfr_slice)
    deallocate(dwfs_slice, dwfr_slice)
    deallocate(kx2, ky2, tapx, tapy)

    deallocate(plan_2dsF,plan_2dsI)
    deallocate(plan_2drF,plan_2drI)
    deallocate(plan_2dtF,plan_2dtI)

    deallocate(index_mxs,index_mxr)
    deallocate(index_mys,index_myr)

    deallocate(cwrk1d) !!.. ifd or ffd
    deallocate(cwrk3d) !!.. ssf-epi
    deallocate(iwrk2d) !!.. ssf-epi
    deallocate(rwrk1d)

    stat = 0
  end function
  !!-----------------------------------------------------------------------------------------
  integer function wemva_kerpix_oper(adj, add, dslw_, dimg_) result(stat)
    logical           :: adj, add
    character(len=*)  :: dslw_, dimg_
    type(axis)        :: aws
    integer           :: izs, ish
    integer           :: ogrid_nx_bak
    integer           :: ogrid_ny_bak
    integer           :: ogrid_nz_bak

    if (adj) then
       call wemva_histo_dslw(vgrid, dslw_)
    else
       call wemva_histo_imag(igrid, dimg_)
    end if

    if (.not.add) then
       if (adj) then
          write(0,*)"nullifing dslw_..."
          call wemva_fio_dslw_init(vgrid, dslw_)
       else
          write(0,*)"nullifing dimg_..."
          call wemva_fio_imag_init(igrid, dimg_)
       end if
    end if

!   loop over shots
!   ---------------
    do ish=1, dgrid%nshg

       ogrid_nx_bak = ogrid%ax%n; ogrid_ny_bak = ogrid%ay%n; ogrid_nz_bak = ogrid%az%n

       call wemva_glocal(pgrid, ogrid, vgrid, dgrid, igrid, ish)
       if (igrid%ax%n>1.and.ogrid%ax%n == 0) cycle
       if (igrid%ay%n>1.and.ogrid%ay%n == 0) cycle
       if (ogrid%ax%n/=ogrid_nx_bak .or. ogrid%ay%n/=ogrid_ny_bak .or. ogrid%az%n/=ogrid_nz_bak) then
          write(0,*)"arrays have been reallocated"
          deallocate(dimg); 
          deallocate(index_mxs,index_mxr)
          deallocate(index_mys,index_myr)  
          allocate(dimg(ogrid%ahx%n,ogrid%ahy%n,ogrid%ax%n,ogrid%ay%n,nth))
          allocate(index_mxs(ogrid%ahx%n,ogrid%ax%n),index_mxr(ogrid%ahx%n,ogrid%ax%n)) 
          allocate(index_mys(ogrid%ahy%n,ogrid%ay%n),index_myr(ogrid%ahy%n,ogrid%ay%n))
       end if

       write(0,*)"------------------------------------------------------------------------------"
       write(0,*)"       ish=", ish, "@", pgrid%sx, pgrid%sy, dgrid%nshg
       write(0,*)"------------------------------------------------------------------------------"

       stat = wem_wfdexp_init(kx2, ky2, pgrid%ax%n, pgrid%ay%n, pgrid%ax%d, pgrid%ay%d, pgrid%az%d, &
                              pgrid%pdip, pgrid%nvrf, pgrid%oper, pgrid%trick)

       stat = wem_velref_init(pgrid%nvrf)

       stat = wem_imgcon_init(ogrid, pgrid, dgrid%aw%n, dgrid%nshg, index_mxs, index_mxr, index_mys, index_myr)

       stat = wemva_bornsc_init(kx2, ky2, pgrid%ax%d, pgrid%ay%d, pgrid%az%d, &
                                pgrid%ax%n, pgrid%ay%n, etabsc, nopbsc, typbsc)

       stat = wemva_tompix_oper(adj, add, dslw_, dimg_, ish)

    end do
    stat = 0
  end function

  !!-----------------------------------------------------------------------------------------
  !! compute the background image
  !!-----------------------------------------------------------------------------------------
  integer function wemva_bimpix_oper(add, bslw_, bimg_) result(stat)
    logical          :: add
    character(len=*) :: bslw_, bimg_
    integer          :: ish
    type(axis)       :: aws
    integer          :: ith, iwb, iws, izs
    real             :: ws
    integer          :: ogrid_nx_bak
    integer          :: ogrid_ny_bak
    integer          :: ogrid_nz_bak

    !!.. nullify background image if requested
    if (.not.add) call wemva_fio_imag_init(igrid, bimg_)

    do ish=1, dgrid%nshg

       ogrid_nx_bak = ogrid%ax%n; ogrid_ny_bak = ogrid%ay%n; ogrid_nz_bak = ogrid%az%n
       call wemva_glocal(pgrid, ogrid, vgrid, dgrid, igrid, ish)

       if (igrid%ax%n>1.and.ogrid%ax%n == 0) cycle
       if (igrid%ay%n>1.and.ogrid%ay%n == 0) cycle

       if (ogrid%ax%n/=ogrid_nx_bak .or. ogrid%ay%n/=ogrid_ny_bak .or. ogrid%az%n/=ogrid_nz_bak) then
          write(0,*)"arrays have been reallocated"
          deallocate(dimg); 
          deallocate(index_mxs,index_mxr)
          deallocate(index_mys,index_myr) 
          allocate(dimg(ogrid%ahx%n,ogrid%ahy%n,ogrid%ax%n,ogrid%ay%n,nth))
          allocate(index_mxs(ogrid%ahx%n,ogrid%ax%n),index_mxr(ogrid%ahx%n,ogrid%ax%n)) 
          allocate(index_mys(ogrid%ahy%n,ogrid%ay%n),index_myr(ogrid%ahy%n,ogrid%ay%n))
       end if

       write(0,*)"------------------------------------------------------------------------------"
       write(0,*)"       ish=", ish, "@", pgrid%sx, pgrid%sy, dgrid%nshg
       write(0,*)"------------------------------------------------------------------------------"

       stat = wem_wfdexp_init(kx2, ky2, pgrid%ax%n, pgrid%ay%n, pgrid%ax%d, pgrid%ay%d, pgrid%az%d, &
                              pgrid%pdip, pgrid%nvrf, pgrid%oper, pgrid%trick)

       stat = wem_velref_init(pgrid%nvrf)

       stat = wem_imgcon_init(ogrid, pgrid, dgrid%aw%n, dgrid%nshg, index_mxs, index_mxr, index_mys, index_myr)

       do iwb=1, wblock%nb
          aws%o = wblock%os(iwb); aws%d = wblock%ds; aws%n = wblock%blen(iwb)
          bwfs = 0; bwfr = 0.
          call wemva_fio_csou(csou, dgrid, wblock%indx(iwb), aws%n, ish, ftag_%csou_, "reed")
          call wemva_fio_crec(crec, dgrid, wblock%indx(iwb), aws%n, ish, ftag_%crec_, "reed")
          call wemva_wfs_init(.true., .false., bwfs(:,:,:,1), csou, dgrid, pgrid)
          call wemva_wfr_init(.true., .false., bwfr(:,:,:,1), crec, dgrid, pgrid)

          do izs=1, pgrid%az%n
             call wemva_fio_velo(cvel, vgrid, pgrid, izs, ftag_%bslw_, "reed")
             bvel = 1./real(cvel)
             stat = wem_velref_oper(bvel, bvrf)

             bimg = 0.
             call wemva_fio_imag(bimg(:,:,:,:,1), igrid, ogrid, izs, bimg_, "reed")

             !$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) FIRSTPRIVATE(iws,izs,ith,ws)
             do iws=1, aws%n
                ith = omp_get_thread_num() + 1
                ws  = aws%o + (iws-1)*aws%d
                dwfs_slice(:,:,ith) = bwfs(:,:,iws,izs)
                dwfr_slice(:,:,ith) = bwfr(:,:,iws,izs)
                stat = wem_imgcon_oper(.false., bimg(:,:,:,:,ith), dwfs_slice(:,:,ith), dwfr_slice(:,:,ith), ws, "all")
                stat = wem_wfdexp_oper(dwfs_slice(:,:,ith), bvel, bvrf, ws, -1,                         &
                       iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dsF(ith), plan_2dsI(ith))
                stat = wem_wfdexp_oper(dwfr_slice(:,:,ith), bvel, bvrf, ws, -1,                         &
                       iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2drF(ith), plan_2drI(ith))
                stat = wem_tapers_oper(dwfs_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                stat = wem_tapers_oper(dwfr_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                bwfs(:,:,iws,izs+1) = dwfs_slice(:,:,ith)
                bwfr(:,:,iws,izs+1) = dwfr_slice(:,:,ith)
             end do
             !$OMP END PARALLEL DO
             call wemva_fio_imag(sum(bimg,5), igrid, ogrid, izs, bimg_, "rite")
          end do
       end do
       stat = wem_velref_clos()
       stat = wem_wfdexp_clos()
       stat = wem_imgcon_clos()
    end do
    stat = 0
  end function

  !!-------------------------------------------------------------------------------------------
  integer function wemva_tompix_oper(adj, add, dslw_, dimg_, ish) result(stat)
    logical           :: adj
    logical           :: add
    character(len=*)  :: dslw_
    character(len=*)  :: dimg_
    integer           :: ish

    type(axis)        :: aws
    integer           :: ith, iwb, iws, izs, jx, jz
    real              :: ws

    call sseek(dimg_, 0, 0)
    call sseek(dslw_, 0, 0)

    if (adj) then     !!.. dS = F' dI

       do iwb=1, wblock%nb
          aws%o = wblock%os(iwb); aws%d = wblock%ds; aws%n = wblock%blen(iwb)
          bwfs = 0.; bwfr = 0.
          dwfs = 0.; dwfr = 0.
          call wemva_fio_csou(csou, dgrid, wblock%indx(iwb), aws%n, ish, ftag_%csou_, "reed")
          call wemva_fio_crec(crec, dgrid, wblock%indx(iwb), aws%n, ish, ftag_%crec_, "reed")
          call wemva_wfs_init(.true., .false., bwfs(:,:,:,1), csou, dgrid, pgrid)
          call wemva_wfr_init(.true., .false., bwfr(:,:,:,1), crec, dgrid, pgrid)
          
          do izs=1, pgrid%az%n-1
             call wemva_fio_velo(cvel, vgrid, pgrid, izs, ftag_%bslw_, "reed")
             bvel = 1./real(cvel)
             stat = wem_velref_oper(bvel, bvrf)
             dwfs_slice = 0.; dwfr_slice = 0. 
             !$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) FIRSTPRIVATE(iws,izs,ith,ws)
             do iws=1, aws%n
                ith = omp_get_thread_num() + 1
                ws  = aws%o + (iws-1)*aws%d
                dwfs_slice(:,:,ith) = bwfs(:,:,iws,izs)
                dwfr_slice(:,:,ith) = bwfr(:,:,iws,izs)
                stat = wem_wfdexp_oper(dwfs_slice(:,:,ith), bvel, bvrf, ws, -1, iwrk2d(:,:,ith), rwrk1d(:,ith), &
                                       cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dsF(ith), plan_2dsI(ith))
                stat = wem_wfdexp_oper(dwfr_slice(:,:,ith), bvel, bvrf, ws, -1, iwrk2d(:,:,ith), rwrk1d(:,ith), &
                                       cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2drF(ith), plan_2drI(ith))
                stat = wem_tapers_oper(dwfs_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                stat = wem_tapers_oper(dwfr_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                bwfs(:,:,iws,izs+1) = dwfs_slice(:,:,ith)
                bwfr(:,:,iws,izs+1) = dwfr_slice(:,:,ith)
             end do
             !$OMP END PARALLEL DO
          end do

          !!.. dP_i = E_{i-1}(dP_{i-1}+G_{i-1}dS_{i-1})
          do izs=pgrid%az%n-1, 1, -1
             call wemva_fio_velo(cvel, vgrid, pgrid, izs, ftag_%bslw_, "reed")          !!.. bS_{i}
             bvel = 1./real(cvel)
             stat = wem_velref_oper(bvel, bvrf)

             dimg = 0.; dslw = 0.
             call wemva_fio_imag(dimg(:,:,:,:,1), igrid, ogrid, izs+1, dimg_, "reed")   !!.. dI_{i+1}
             call wemva_fio_velo(dslw(:,:,1)    , vgrid, pgrid, izs  , dslw_, "reed")   !!.. dS_i

             do ith=1, nth
                dimg(:,:,:,:,ith) = dimg(:,:,:,:,1)
             end do

             bwfs_slice = 0.; bwfr_slice = 0.
             dwfs_slice = 0.; dwfr_slice = 0.
             wtmp_slice = 0.
             !$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) FIRSTPRIVATE(iws,izs,ith,ws)
             do iws=1, aws%n
                ith = omp_get_thread_num() + 1
                ws  = aws%o + (iws-1)*aws%d

!               source part
!               -----------
                bwfr_slice(:,:,ith) = bwfr(:,:,iws,izs+1)    !!.. bU_{i+1}
                dwfs_slice(:,:,ith) = dwfs(:,:,iws,izs+1)    !!.. dP_{i+1}
                stat = wem_imgcon_oper(.true., dimg(:,:,:,:,ith), dwfs_slice(:,:,ith), bwfr_slice(:,:,ith), ws, "sou")   !!.. dP_{i+1} = dP_{i+1} + bU_{i+1}dI_{i+1}
                stat = wem_tapers_oper(dwfs_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)          
                stat = wem_wfdexp_oper(dwfs_slice(:,:,ith), bvel, bvrf, ws, +1, iwrk2d(:,:,ith), rwrk1d(:,ith), &        !!.. dP_i = E_i' dP_{i+1}
                                       cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dsF(ith), plan_2dsI(ith))
                stat = wemva_bornsc_oper(bwfs_slice(:,:,ith), bwfs(:,:,iws,izs), wtmp_slice(:,:,ith), &                  !!.. G_i
                                         bvel, ws, -1, plan_2dtF(ith), plan_2dtI(ith))
                dslw(:,:,ith) = dslw(:,:,ith) + conjg(bwfs_slice(:,:,ith))*dwfs_slice(:,:,ith)                           !!.. dS_i = G_i' dP_i
                dwfs(:,:,iws,izs) = dwfs_slice(:,:,ith)                                               

!               receiver part
!               -------------
                bwfs_slice(:,:,ith) = bwfs(:,:,iws,izs+1)
                dwfr_slice(:,:,ith) = dwfr(:,:,iws,izs+1)
                stat = wem_imgcon_oper(.true., dimg(:,:,:,:,ith), bwfs_slice(:,:,ith), dwfr_slice(:,:,ith), ws, "rec")
                stat = wem_tapers_oper(dwfr_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                stat = wem_wfdexp_oper(dwfr_slice(:,:,ith), bvel, bvrf, ws, +1, iwrk2d(:,:,ith), rwrk1d(:,ith), &
                                       cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2drF(ith), plan_2drI(ith))
                stat = wemva_bornsc_oper(bwfr_slice(:,:,ith), bwfr(:,:,iws,izs), wtmp_slice(:,:,ith), &
                                         bvel, ws, -1, plan_2dtF(ith), plan_2dtI(ith))
                dslw(:,:,ith) = dslw(:,:,ith) + conjg(bwfr_slice(:,:,ith))*dwfr_slice(:,:,ith)
                dwfr(:,:,iws,izs) = dwfr_slice(:,:,ith) 

             end do
             !$OMP END PARALLEL DO
             call wemva_fio_velo(sum(dslw,3), vgrid, pgrid, izs, dslw_, "rite")
          end do
       end do


    else              !!.. dI = F  dS
       
       do iwb=1, wblock%nb
          aws%o = wblock%os(iwb); aws%d = wblock%ds; aws%n = wblock%blen(iwb)
          bwfs = 0.; bwfr = 0.
          dwfs = 0.; dwfr = 0.
          call wemva_fio_csou(csou, dgrid, wblock%indx(iwb), aws%n, ish, ftag_%csou_, "reed")
          call wemva_fio_crec(crec, dgrid, wblock%indx(iwb), aws%n, ish, ftag_%crec_, "reed")
          call wemva_wfs_init(.true., .false., bwfs(:,:,:,1), csou, dgrid, pgrid)
          call wemva_wfr_init(.true., .false., bwfr(:,:,:,1), crec, dgrid, pgrid)

          do izs=1, pgrid%az%n-1
             call wemva_fio_velo(cvel, vgrid, pgrid, izs, ftag_%bslw_, "reed")
             bvel = 1./real(cvel)
             stat = wem_velref_oper(bvel, bvrf)

             dwfs_slice = 0.; dwfr_slice = 0.
             !$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) FIRSTPRIVATE(iws,izs,ith,ws)
             do iws=1, aws%n
                ith = omp_get_thread_num() + 1
                ws  = aws%o + (iws-1)*aws%d
                dwfs_slice(:,:,ith) = bwfs(:,:,iws,izs)
                dwfr_slice(:,:,ith) = bwfr(:,:,iws,izs)
                stat = wem_wfdexp_oper(dwfs_slice(:,:,ith), bvel, bvrf, ws, -1, iwrk2d(:,:,ith), rwrk1d(:,ith), &
                                       cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dsF(ith), plan_2dsI(ith))
                stat = wem_wfdexp_oper(dwfr_slice(:,:,ith), bvel, bvrf, ws, -1, iwrk2d(:,:,ith), rwrk1d(:,ith), &
                                       cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2drF(ith), plan_2drI(ith))
                stat = wem_tapers_oper(dwfs_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                stat = wem_tapers_oper(dwfr_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                bwfs(:,:,iws,izs+1) = dwfs_slice(:,:,ith)
                bwfr(:,:,iws,izs+1) = dwfr_slice(:,:,ith)
             end do
             !$OMP END PARALLEL DO
          end do

          !!.. dP_i = E_{i-1}(dP_{i-1}+G_{i-1}dS_{i-1})
          do izs=2, pgrid%az%n
             call wemva_fio_velo(cvel, vgrid, pgrid, izs-1, ftag_%bslw_, "reed")          !!.. bS_{i-1}
             bvel = 1./real(cvel)
             stat = wem_velref_oper(bvel, bvrf)

             dimg = 0.; dslw = 0.
             call wemva_fio_imag(dimg(:,:,:,:,1), igrid, ogrid, izs  , dimg_, "reed")        !!.. dI_i
             call wemva_fio_velo(dslw(:,:,1)    , vgrid, pgrid, izs-1, dslw_, "reed")      !!.. dS_{i-1}

             do ith=1, nth
                dslw(:,:,ith) = dslw(:,:,1)
             end do

             bwfs_slice = 0.; bwfr_slice = 0.
             dwfs_slice = 0.; dwfr_slice = 0.
             wtmp_slice = 0.
             !$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) FIRSTPRIVATE(iws,izs,ith,ws)
             do iws=1, aws%n
                ith = omp_get_thread_num() + 1
                ws  = aws%o + (iws-1)*aws%d

!               source part
!               -----------
                bwfr_slice(:,:,ith) = bwfr(:,:,iws,izs  )    !!.. bU_i
                dwfs_slice(:,:,ith) = dwfs(:,:,iws,izs-1)    !!.. dD_{i-1}
                stat = wemva_bornsc_oper(bwfs_slice(:,:,ith), bwfs(:,:,iws,izs-1), wtmp_slice(:,:,ith), & !!.. G_{i-1}
                                         bvel, ws, -1, plan_2dtF(ith), plan_2dtI(ith))
                dwfs_slice(:,:,ith) = dwfs_slice(:,:,ith) + bwfs_slice(:,:,ith)*dslw(:,:,ith)             !!.. dD_{i-1} = dD_{i-1} + G_{i-1}dS_{i-1}
                stat = wem_wfdexp_oper(dwfs_slice(:,:,ith), bvel, bvrf, ws, -1, iwrk2d(:,:,ith), rwrk1d(:,ith), &   !!.. dD_i = E_{i-1} dD_{i-1}
                                       cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dsF(ith), plan_2dsI(ith))
                stat = wem_tapers_oper(dwfs_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                stat = wem_imgcon_oper(.false., dimg(:,:,:,:,ith), dwfs_slice(:,:,ith), bwfr_slice(:,:,ith), ws, "sou") !!.. dI_i = dD_i' bU_i
                dwfs(:,:,iws,izs) = dwfs_slice(:,:,ith)


!               receiver part
!               -------------
                bwfs_slice(:,:,ith) = bwfs(:,:,iws,izs  )
                dwfr_slice(:,:,ith) = dwfr(:,:,iws,izs-1)
                stat = wemva_bornsc_oper(bwfr_slice(:,:,ith), bwfr(:,:,iws,izs-1), wtmp_slice(:,:,ith), &
                                         bvel, ws, -1, plan_2dtF(ith), plan_2dtI(ith))
                dwfr_slice(:,:,ith) = dwfr_slice(:,:,ith) + bwfr_slice(:,:,ith)*dslw(:,:,ith)
                stat = wem_wfdexp_oper(dwfr_slice(:,:,ith), bvel, bvrf, ws, -1, iwrk2d(:,:,ith), rwrk1d(:,ith), & 
                                       cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2drF(ith), plan_2drI(ith))
                stat = wem_tapers_oper(dwfr_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                stat = wem_imgcon_oper(.false., dimg(:,:,:,:,ith), bwfs_slice(:,:,ith), dwfr_slice(:,:,ith), ws, "rec")
                dwfr(:,:,iws,izs) = dwfr_slice(:,:,ith)

             end do
             !$OMP END PARALLEL DO
             call wemva_fio_imag(sum(dimg,5), igrid, ogrid, izs, dimg_, "rite")
          end do

       end do

    end if
  end function

end module
