module wemva_kernel_mod2
  use  base_struct_mod
  use  base_blkcal_mod
  use  base_ioutil_mod
  use  wem_velref_mod
  use  wem_wfdexp_mod
  use  wem_tapers_mod
  use  wem_fftwv3_mod
  use  wemva_glocal_mod
  use  wemva_fileio_mod
  use  wemva_imgcon_mod
  use  wemva_phsscr_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

  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

  real     , dimension(:,:,:)  , allocatable, private :: bvel
  real     , dimension(:,:,:)  , allocatable, private :: bimg
  real     , dimension(:,:,:)  , allocatable, private :: dslw
  real     , 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

  character(len=256)                        , private :: csou_
  character(len=256)                        , private :: crec_
  character(len=256)                        , private :: hsou_
  character(len=256)                        , private :: bvel_
  character(len=256)                        , private :: bimg_
  character(len=256)                        , private :: bwfs_
  character(len=256)                        , private :: bwfr_
  character(len=256)                        , private :: dslw_
  character(len=256)                        , private :: dimg_

  integer                                   , private :: ompchunk, nth, nwstep
  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

  logical                                   , private :: adj, add
  logical                                   , private :: wntbim
  real   ,                         parameter, private :: PI = 3.14159265359
contains

  integer function wemva_prestk_init(adj, add) result(stat)

    call wemva_param_init(dgrid, pgrid, igrid, wblock, crec_, csou_, hsou_, &
                          bvel_, bimg_, bwfs_, bwfr_, dslw_, dimg_)
 
!    call from_param("adj"   , adj   , .false.)
!    call from_param("add"   , add   , .false.)
    call from_param("wntbim", wntbim, .true. )
    call from_param("nwstep", nwstep, 100    )

    call wemva_histo_data(dgrid, vgrid, igrid, pgrid, crec_, csou_, bvel_, hsou_)

    if (wntbim) call wemva_histo_imag(igrid, bimg_)

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

    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, 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,pgrid%az%n,wblock%ns))
    allocate(bwfr(pgrid%ax%n,pgrid%ay%n,pgrid%az%n,wblock%ns))
    allocate(dwfs(pgrid%ax%n,pgrid%ay%n,pgrid%az%n,wblock%ns))
    allocate(dwfr(pgrid%ax%n,pgrid%ay%n,pgrid%az%n,wblock%ns))
    allocate(bvel(pgrid%ax%n,pgrid%ay%n,pgrid%az%n))
    allocate(bvrf(pgrid%nvrf,pgrid%az%n))
    allocate(dslw(pgrid%ax%n,pgrid%ay%n,pgrid%az%n))
    if (wntbim) allocate(bimg(ogrid%ax%n,ogrid%ay%n,ogrid%az%n))
    allocate(dimg(ogrid%ax%n,ogrid%ay%n,ogrid%az%n))
    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(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(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)

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

    if (wntbim) then
       call wemva_fio_imag_init(igrid, bimg_)
    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

    stat = 0
  end function
  !!-----------------------------------------------------------------------------------------
  integer function wemva_prestk_clos() result(stat)
    stat = 0
  end function
  !!-----------------------------------------------------------------------------------------
  integer function wemva_prestk_oper() result(stat)
    type(axis)        :: aws
    integer           :: izs, ish
    integer           :: ogrid_nx_bak
    integer           :: ogrid_ny_bak
    integer           :: ogrid_nz_bak

    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"
          if (wntbim) then
             deallocate(bimg); allocate(bimg(ogrid%ax%n,ogrid%ay%n,ogrid%az%n)); bimg = 0.
          end if
          deallocate(dimg); allocate(dimg(ogrid%ax%n,ogrid%ay%n,ogrid%az%n)); dimg = 0.
       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)
       call wemva_fio_velo(bvel, vgrid, pgrid, 1, pgrid%az%n, bvel_, "reed")
       do izs=1, pgrid%az%n
          stat = wem_velref_oper(bvel(:,:,izs), bvrf(:,izs))
       end do
       stat = wemva_imgcon_init(igrid, ogrid, pgrid)
       stat = wemva_phsscr_init(pgrid%ax%d, pgrid%ay%d, pgrid%az%d, pgrid%ax%n, pgrid%ay%n)
       stat = wemva_bgwfld_oper(ish)
       stat = wemva_tomgph_oper(adj, add, ish)

    end do
    stat = 0
  end function

  !!-----------------------------------------------------------------------------------------
  integer function wemva_bgwfld_oper(ish) result(stat)
    integer    :: ish
    type(axis) :: aws
    integer    :: ith, iwb, iws, izs
    real       :: ws

    call sseek(bwfs_, 0, 0)
    call sseek(bwfr_, 0, 0)
    if (wntbim) call sseek(bimg_, 0, 0)

    call wemva_histo_wfld(pgrid, dgrid%aw, bwfs_)
    call wemva_histo_wfld(pgrid, dgrid%aw, bwfr_)
    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, csou_, "reed")
       call wemva_fio_crec(crec, dgrid, wblock%indx(iwb), aws%n, ish, crec_, "reed")
       call wemva_wfs_init(bwfs, csou, dgrid, pgrid)
       call wemva_wfr_init(bwfr, crec, dgrid, pgrid)
       iwrk2d = 0; rwrk1d = 0.; cwrk1d = 0.; cwrk3d = 0.

       !$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) PRIVATE(iws,izs,ith,ws)
       do iws=1, aws%n
          ith  = omp_get_thread_num() + 1
          ws   = aws%o + (iws-1)*aws%d
          if (iws/nwstep*nwstep==iws) write(0,*)"iws=", iws, ws, ith

          dwfs_slice(:,:,ith) = bwfs(:,:,1,iws)
          dwfr_slice(:,:,ith) = bwfr(:,:,1,iws)

          do izs=2, pgrid%nzstep, +1
             stat = wem_wfdexp_oper(dwfs_slice(:,:,ith), bvel(:,:,izs-1), bvrf(:,izs-1), ws, +1,                     &
                    iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dsF(ith), plan_2dsI(ith))
             stat = wem_wfdexp_oper(dwfr_slice(:,:,ith), bvel(:,:,izs-1), bvrf(:,izs-1), 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(:,:,izs,iws) = dwfs_slice(:,:,ith)
             bwfr(:,:,izs,iws) = dwfr_slice(:,:,ith)
          end do

          if (wntbim) then
             !$OMP CRITICAL
             stat = wemva_imgcon_oper(.false., bimg, bwfs(:,:,:,iws), bwfr(:,:,:,iws), ws, "all", bimg_)
             !$OMP END CRITICAL
          end if

       end do
       !$OMP END PARALLEL DO
       call sep_write(bwfs(:,:,:,1:aws%n), bwfs_)
       call sep_write(bwfr(:,:,:,1:aws%n), bwfr_)
    end do
    stat = 0
  end function
  !!-------------------------------------------------------------------------------------------
  integer function wemva_tomgph_oper(adj, add, ish) result(stat)
    logical                                           :: adj
    logical                                           :: add
    !real, dimension(pgrid%ax%n,pgrid%ay%n,pgrid%az%n) :: dslw
    !real, dimension(ogrid%ax%n,ogrid%ay%n,ogrid%az%n) :: dimg
    integer                                           :: ish

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


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

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

       call wemva_fio_velo(dslw, vgrid, pgrid, 1, pgrid%az%n, dslw_, "reed")

       do iwb=1, wblock%nb
          aws%o = wblock%os(iwb); aws%d = wblock%ds; aws%n = wblock%blen(iwb)
          bwfs = 0.; dwfs = 0.
          bwfr = 0.; dwfr = 0.
          call sep_read(bwfs(:,:,:,1:aws%n), bwfs_)
          call sep_read(bwfr(:,:,:,1:aws%n), bwfr_)
          iwrk2d = 0; rwrk1d = 0.; cwrk1d = 0.; cwrk3d = 0.
          !$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) PRIVATE(iws,izs,ith,ws)
          do iws=1, aws%n
             ith  = omp_get_thread_num() + 1
             ws   = aws%o + (iws-1)*aws%d
             if (iws/nwstep*nwstep==iws) write(0,*)"iws=", iws, ws, ith
             !$OMP CRITICAL
             stat = wemva_imgcon_oper(adj, dimg, dwfs(:,:,:,iws), bwfr(:,:,:,iws), ws, "sou", dimg_)
             stat = wemva_imgcon_oper(adj, dimg, bwfs(:,:,:,iws), dwfr(:,:,:,iws), ws, "rec", dimg_)
             !$OMP END CRITICAL
             bwfs_slice(:,:,ith) = 0.; dwfs_slice(:,:,ith) = 0.
             bwfr_slice(:,:,ith) = 0.; dwfr_slice(:,:,ith) = 0.
             do izs=pgrid%nzstep, 2, -1
                dwfs_slice(:,:,ith) = dwfs_slice(:,:,ith) + dwfs(:,:,izs,iws) 
                stat = wem_tapers_oper(dwfs_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                stat = wem_wfdexp_oper(dwfs_slice(:,:,ith), bvel(:,:,izs-1), bvrf(:,izs-1), ws, -1,                     &
                       iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dsF(ith), plan_2dsI(ith))
                stat = wemva_phsscr_oper(bwfs_slice(:,:,ith), bwfs(:,:,izs-1,iws), bvel(:,:,izs-1), ws,+1)
                dslw(:,:,izs-1) = dslw(:,:,izs-1) + conjg(bwfs_slice(:,:,ith))*dwfs_slice(:,:,ith) !!.. bug danger

                dwfr_slice(:,:,ith) = dwfr_slice(:,:,ith) + dwfr(:,:,izs,iws)
                stat = wem_tapers_oper(dwfr_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                stat = wem_wfdexp_oper(dwfr_slice(:,:,ith), bvel(:,:,izs-1), bvrf(:,izs-1), ws, +1,                     &
                       iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2drF(ith), plan_2drI(ith))
                stat = wemva_phsscr_oper(bwfr_slice(:,:,ith), bwfr(:,:,izs-1,iws), bvel(:,:,izs-1), ws, -1)
                dslw(:,:,izs-1) = dslw(:,:,izs-1) + conjg(bwfr_slice(:,:,ith))*dwfr_slice(:,:,ith) !!.. bug danger
             end do

          end do
          !$OMP END PARALLEL DO
       end do

       call wemva_fio_velo(dslw, vgrid, pgrid, 1, pgrid%az%n, dslw_, "rite")

    else              !!.. dI = F  dS
       
       call wemva_fio_velo(dslw, vgrid, pgrid, 1, pgrid%az%n, dslw_, "reed")
 
       do iwb=1, wblock%nb
          aws%o = wblock%os(iwb); aws%d = wblock%ds; aws%n = wblock%blen(iwb)
          bwfs = 0.; dwfs = 0.
          bwfr = 0.; dwfr = 0.
          call sep_read(bwfs(:,:,:,1:aws%n), bwfs_)
          call sep_read(bwfr(:,:,:,1:aws%n), bwfr_)
          iwrk2d = 0; rwrk1d = 0.; cwrk1d = 0.; cwrk3d = 0.
          !$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) PRIVATE(iws,izs,ith,ws)
          do iws=1, aws%n
             ith  = omp_get_thread_num() + 1
             ws   = aws%o + (iws-1)*aws%d
             if (iws/nwstep*nwstep==iws) write(0,*)"iws=", iws, ws, ith
             bwfs_slice(:,:,ith) = 0.; dwfs_slice(:,:,ith) = 0.
             bwfr_slice(:,:,ith) = 0.; dwfr_slice(:,:,ith) = 0.
             !!.. dP_i = E_{i-1}(dP_{i-1}+G_{i-1}dS_{i-1})
             do izs=2, pgrid%az%n
                stat = wemva_phsscr_oper(bwfs_slice(:,:,ith), bwfs(:,:,izs-1,iws), bvel(:,:,izs-1), ws, +1)                !!.. G_{i-1}
                dwfs_slice(:,:,ith) = dwfs_slice(:,:,ith) + bwfs_slice(:,:,ith)*dslw(:,:,izs-1)                            !!.. dP_{i-1}+G_{i-1}dS_{i-1}
                stat = wem_wfdexp_oper(dwfs_slice(:,:,ith), bvel(:,:,izs-1), bvrf(:,izs-1), ws, +1,                     &
                       iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dsF(ith), plan_2dsI(ith))   !!.. E_{i-1}(dP_{i-1}+G_{i-1}dS_{i-1}) 
                stat = wem_tapers_oper(dwfs_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                dwfs(:,:,izs,iws) = dwfs_slice(:,:,ith)

                stat = wemva_phsscr_oper(bwfr_slice(:,:,ith), bwfr(:,:,izs-1,iws), bvel(:,:,izs-1), ws, -1)
                dwfr_slice(:,:,ith) = dwfr_slice(:,:,ith) + bwfr_slice(:,:,ith)*dslw(:,:,izs-1)
                stat = wem_wfdexp_oper(dwfr_slice(:,:,ith), bvel(:,:,izs-1), bvrf(:,izs-1), 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)
                dwfr(:,:,izs,iws) = dwfr_slice(:,:,ith)
             end do
             !$OMP CRITICAL
             stat = wemva_imgcon_oper(adj, dimg, dwfs(:,:,:,iws), bwfr(:,:,:,iws), ws, "sou", dimg_)
             stat = wemva_imgcon_oper(adj, dimg, bwfs(:,:,:,iws), dwfr(:,:,:,iws), ws, "rec", dimg_)
             !$OMP END CRITICAL
          end do
          !$OMP END PARALLEL DO
       end do

    end if
  end function


  subroutine wemva_arrays_rset()
    csou = 0.; crec = 0.
    bwfs = 0.; bwfr = 0.
    dwfs = 0.; bwfr = 0.
    bvel = 0.; bvrf = 0.
    dslw = 0.; dimg = 0.
    if (wntbim) bimg = 0.
    bwfs_slice = 0.; bwfr_slice = 0.
    dwfs_slice = 0.; dwfr_slice = 0.
    cwrk1d = 0.; cwrk3d = 0.
    iwrk2d = 0 ; rwrk1d = 0.
    return
  end subroutine

end module
