module mva_postom_mod
  use  sep
  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  mva_phsscr_mod  
  implicit none

  type(axis)                                , private :: ax, ay, az, aw, aws
  type(axis)                                , private :: arx, ary, avx, avy
  type(block)                               , private :: wblock

  complex  , dimension(:,:,:,:), allocatable, private :: bP
  complex  , dimension(:,:,:)  , allocatable, private :: dU_slice
  complex  , dimension(:,:,:)  , allocatable, private :: dP_slice
  complex  , dimension(:,:,:)  , allocatable, private :: D
  real     , dimension(:,:,:)  , allocatable, private :: bI
  real     , dimension(:,:,:)  , allocatable, private :: dI
  real     , dimension(:,:,:)  , allocatable, private :: bS
  real     , dimension(:,:,:)  , allocatable, private :: dS
  real     , dimension(:,:,:)  , allocatable, private :: bV
  real     , dimension(:,:)    , allocatable, private :: bVr

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

  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_2ddPF, plan_2ddPI

  character(len=5)                          , private :: oper
  integer                                   , private :: nxpad_beg, nxpad_end
  integer                                   , private :: nypad_beg, nypad_end
  integer                                   , private :: pdip, nvrf, ntpx, ntpy
  integer                                   , private :: ompchunk, nth, ith
  real                                      , private :: trick

  logical                                   , private :: adj, radian

  character(len=256)                        , private :: bS_, dS_, D_, bI_, dI_, bP_

  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


contains

  integer function mva_pospar_init() result(stat)
    integer :: ix, iy, iz
    real    :: rx_min, rx_max
    real    :: ry_min, ry_max
    real    :: px_min, px_max
    real    :: py_min, py_max
    real    :: PI

    PI = 3.14159265359

    bS_ = "bS"
    dS_ = "dS"
    D_  = "D"
    bI_ = "bI"
    dI_ = "dI"
    bP_ = "bP"

    ompchunk = 1
    nth      = 1
    ith      = 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


    call from_param("pdip"     , pdip     ,      45)
    call from_param("nvrf"     , nvrf     ,       1)
    call from_param("trick"    , trick    ,    0.14)
    call from_param("ntpx"     , ntpx     ,       1)
    call from_param("ntpy"     , ntpy     ,       1)
    call from_param("oper"     , oper     , "ssfpi")
    call from_param("nws"      , wblock%ns,       1)
    call from_param("adj"      , adj      , .false.)
    call from_param("radian"   , radian   ,  .true.)

    call from_param("nxpad_beg", nxpad_beg,       0)
    call from_param("nxpad_end", nxpad_end,       0)
    call from_param("nypad_beg", nypad_beg,       0)
    call from_param("nypad_end", nypad_end,       0)

    !!.. background slowness
    call auxin(bS_)
    call hinaxis(bS_, avx, 1); call hinaxis(bS_, avy, 2); call hinaxis(bS_, az , 3)

    !!.. zero-offset data
    call auxin( D_)
    call hinaxis( D_, arx, 1); call hinaxis( D_, ary, 2); call hinaxis( D_, aw , 3)

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

    if (arx%n/=avx%n .or. ary%n/=avy%n) call seperr("ERROR")

    rx_min = arx%o; rx_max = arx%o + (arx%n-1)*arx%d
    ry_min = ary%o; ry_max = ary%o + (ary%n-1)*ary%d

    px_min = rx_min - nxpad_beg*arx%d; px_max = rx_max + nxpad_end*arx%d
    py_min = ry_min - nypad_beg*ary%d; py_max = ry_max + nypad_end*ary%d

    ax%o = px_min; ax%d = arx%d; ax%n = 1.5 + (px_max-px_min)/ax%d
    ay%o = py_min; ay%d = ary%d; ay%n = 1.5 + (py_max-py_min)/ay%d

    !!.. background wavefield
    call auxinout(bP_)
    call houaxis(bP_, ax , 1); call houaxis(bP_, ay , 2); call houaxis(bP_, az, 3); call houaxis(bP_, aw, 4)
    call auxputch("esize", "i", 8, bP_)

    !!.. background image
    call auxinout(bI_)
    call houaxis(bI_, arx, 1); call houaxis(bI_, ary, 2); call houaxis(bI_, az, 3)

    call base_blkcal(wblock, aw)

    call sep_read(bS(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,:), bS_)

    !! expand domain in x direction
    do ix=1, nxpad_beg; bS(       ix,:,:) = bS(   1+nxpad_beg,:,:); end do
    do ix=1, nxpad_end; bS(ax%n-ix+1,:,:) = bS(ax%n-nxpad_end,:,:); end do

    !! expand domain in y direction
    do iy=1, nypad_beg; bS(:,       iy,:) = bS(:,   1+nypad_beg,:); end do
    do iy=1, nypad_end; bS(:,ay%n-iy+1,:) = bS(:,ay%n-nypad_end,:); end do

    !!.. postack, so twice slowness
    bS = bS * 2.
    bV = 1./bS

    !!.. get the reference velocity for each depth level
    stat = wem_velref_init(nvrf)
    do iz=1, az%n
       stat = wem_velref_oper(bV(:,:,iz), bVr(:,iz))
    end do

    stat = wem_wfdexp_init(kx2, ky2, ax%n, ay%n, ax%d, ay%d, az%d, pdip, nvrf, oper, trick)
    stat = wem_tapers_init(tapx, tapy, ntpx, ntpy)
    stat = mva_phsscr_init(ax%d, ay%d, az%d, ax%n, ay%n)
    
    call k2_map(kx2, ax%n, ax%d)
    call k2_map(ky2, ay%n, ay%d)
    call wem_fftw_init(plan_2ddPF, plan_2ddPI, dP_slice, ax%n, ay%n, nth)

    stat = 0
  end function

  integer function posmva_arrays_aloc() result(stat)
    allocate(bI(ax%n,ay%n,az%n))
    allocate(bS(ax%n,ay%n,az%n))
    allocate(bV(ax%n,ay%n,az%n))
    allocate(bVr(nvrf,az%n))
    allocate(D(ax%n,ay%n,wblock%ns))
    allocate(bP(ax%n,ay%n,az%n,wblock%ns))
    allocate(dP_slice(ax%n,ay%n,nth))
    allocate(kx2(ax%n),ky2(ay%n))
    allocate(tapx(ntpx),tapy(ntpy))

    allocate(plan_2ddPF(nth),plan_2ddPI(nth))

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

  integer function posmva_wfdexp_init() result(stat)
  end function

  integer function mva_posbkg_oper() result(stat)
    integer :: iwb, iws, iz, ith
    real    :: ws

    ith = 1

    call sseek(bP_, 0, 0)
    call sseek(bI_, 0, 0)

    do iwb=1, wblock%nb
       aws%o = wblock%os(iwb); aws%d = wblock%ds; aws%n = wblock%blen(iwb)
       D = 0.
       call sep_read(D(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,1:aws%n), D_)
       dP_slice = 0.; iwrk2d = 0; rwrk1d = 0.; cwrk1d = 0.; cwrk3d = 0.
!$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) PRIVATE(iws,iz,ith,ws)
       do iws=1, aws%n
          ith  = omp_get_thread_num() + 1
          ws   = aws%o + (iws-1)*aws%d
          if (iws/10*10==iws) write(0,*)"iws=", iws, ws, ith
          do iz=1, az%n
             stat = wem_wfdexp_oper(D(:,:,iws), bV(:,:,iz), bVr(:,iz), ws, -1,                          &
                    iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2ddPF(ith), plan_2ddPI(ith))
             stat = wem_tapers_oper(D(:,:,iws), ax%n, ay%n)
             bP(:,:,iz,iws) = D(:,:,iws)
             bI(:,:,iz) = bI(:,:,iz) + D(:,:,iws)/aw%n
          end do
       end do
!$OMP END PARALLEL DO
       call sep_write(bP(:,:,:,1:aws%n), bP_)
       call sep_write(bI(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,:), bI_)
    end do
    stat = 0
  end function

  integer function mva_postom_oper(adj, add, dS, dI, bkg) result(stat)
    logical                :: adj, add, bkg
    real, dimension(:,:,:) :: dS
    real, dimension(:,:,:) :: dI
    integer                :: iwb, iws, iz
    real                   :: ws

    if (bkg) then     !!.. compute the background wavefield
       call auxout(bP_); call sseek(bP_, 0, 0)
       call auxout(bI_); call sseek(bI_, 0, 0)

       do iwb=1, wblock%nb
          aws%o = wblock%os(iwb); aws%d = wblock%ds; aws%n = wblock%blen(iwb)
          call sep_read(D(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,1:aws%n), D_)
          dP_slice = 0.; iwrk2d = 0; rwrk1d = 0.; cwrk1d = 0.; cwrk3d = 0.
!$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) PRIVATE(iws,iz,ith,ws)
          do iws=1, aws%n
             ith  = omp_get_thread_num() + 1
             ws   = aws%o + (iws-1)*aws%d
             if (iws/10*10==iws) write(0,*)"iws=", iws, ws, ith
             dP_slice(:,:,ith) = D(:,:,iws)
             do iz=1, az%n
                stat = wem_wfdexp_oper(dP_slice(:,:,ith), bV(:,:,iz), bVr(:,iz), ws, -1,                          &
                       iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2ddPF(ith), plan_2ddPI(ith))
                stat = wem_tapers_oper(dP_slice(:,:,ith), ax%n, ay%n)
                bP(:,:,iz,iws) = dP_slice(:,:,ith)
                bI(:,:,iz) = bI(:,:,iz) + dP_slice(:,:,ith)/aw%n
             end do
          end do
!$OMP END PARALLEL DO
          call sep_write(bP(:,:,:,1:aws%n), bP_)
          call sep_write(bI(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,:), bI_)
       end do
    end if


    if (.not.add) then
       if (adj) then
          dS = 0.
       else
          dI = 0.
       end if
    end if

    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)
          call sep_read(bP(:,:,:,1:aws%n), bP_)
          iwrk2d = 0; rwrk1d = 0.; cwrk1d = 0.; cwrk3d = 0.
!$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) PRIVATE(iws,iz,ith,ws)
          do iws=1, aws%n
             ith  = omp_get_thread_num() + 1
             ws   = aws%o + (iws-1)*aws%d
             if (iws/10*10==iws) write(0,*)"iws=", iws, ws, ith
             dP_slice(:,:,ith) = 0.
             dU_slice(:,:,ith) = 0.
             do iz=az%n, 1, -1
                stat = mva_phsscr_oper(dU_slice(:,:,ith), bP(:,:,iz,iws), bS(:,:,iz), ws)
                dS(:,:,iz) = dS(:,:,iz) + conjg(dU_slice(:,:,ith))*dP_slice(:,:,ith)/aw%n
                dP_slice(:,:,ith) = dP_slice(:,:,ith) + dI(:,:,iz)
                stat = wem_tapers_oper(dP_slice(:,:,ith), ax%n, ay%n)
                stat = wem_wfdexp_oper(dP_slice(:,:,ith), bV(:,:,iz), bVr(:,iz), ws, +1,                          &
                       iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2ddPF(ith), plan_2ddPI(ith))
                
             end do
          end do
!$OMP END PARALLEL 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)
          call sep_read(bP(:,:,:,1:aws%n), bP_)
          iwrk2d = 0; rwrk1d = 0.; cwrk1d = 0.; cwrk3d = 0.
!$OMP PARALLEL DO SCHEDULE(DYNAMIC,OMPCHUNK) DEFAULT(SHARED) PRIVATE(iws,iz,ith,ws)
          do iws=1, aws%n
             ith  = omp_get_thread_num() + 1
             ws   = aws%o + (iws-1)*aws%d
             if (iws/10*10==iws) write(0,*)"iws=", iws, ws, ith
             dU_slice(:,:,ith) = 0.
             dP_slice(:,:,ith) = 0.
             do iz=1, az%n
                stat = wem_wfdexp_oper(dP_slice(:,:,ith), bV(:,:,iz), bVr(:,iz), ws, -1,                          &
                       iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2ddPF(ith), plan_2ddPI(ith))
                stat = wem_tapers_oper(dP_slice(:,:,ith), ax%n, ay%n)
                dU_slice(:,:,ith) = dU_slice(:,:,ith)*dS(:,:,iz)
                stat = mva_phsscr_oper(dU_slice(:,:,ith), bP(:,:,iz,iws), bS(:,:,iz), ws)
                dP_slice(:,:,ith) = dP_slice(:,:,ith) + dU_slice(:,:,ith)
                dI(:,:,iz)  = dI(:,:,iz)  + dP_slice(:,:,ith)/aw%n
             end do
          end do
!$OMP END PARALLEL DO
       end do
       call sep_write(dI(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,:), dI_)

    end if
  end function

  integer function mva_posbgd_init() result(stat)
  end function
  integer function mva_posbgd_oper() result(stat)
  end function
end module
