module posmva_kernel_mod
  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
  use  posmva_struct_mod
  implicit none

  type(posmva_par)                          , private :: par  

  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

  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

contains

  integer function posmva_params_init(par_out) result(stat)
    type(posmva_par) :: par_out
    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

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

    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

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

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

    call from_param("nwstep"   , nwstep   ,      10)

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

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

    if (.not.par%radian) then
       par%aw%d = par%aw%d*2*PI; par%aw%o = par%aw%o*2*PI
    end if
    if (par%arx%n/=par%avx%n .or. par%ary%n/=par%avy%n) call seperr("ERROR")

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

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

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

    par%jx1 = par%nxpad_beg + 1; par%jx2 = par%nxpad_beg + par%arx%n
    par%jy1 = par%nypad_beg + 1; par%jy2 = par%nypad_beg + par%ary%n

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

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

    call base_blkcal(par%wblock, par%aw)

    call auxinout(par%dS_)
    call auxinout(par%dI_)

    par_out = par

    stat = 0
  end function
  !!--------------------------------------------------------------------------------------
  integer function posmva_tomogph_init(adj) result(stat)
    logical :: adj

    call auxinout(par%dS_)
    call auxinout(par%dI_)

    if (adj)  then
       call houaxis(par%dS_, par%arx, 1); call houaxis(par%dS_, par%ary, 2); call houaxis(par%dS_, par%az, 3)
    else
       call houaxis(par%dI_, par%arx, 1); call houaxis(par%dI_, par%ary, 2); call houaxis(par%dI_, par%az, 3)
    end if
    stat = 0
  end function
  !!---------------------------------------------------------------------------------------
  integer function posmva_arrays_init() result(stat)
    allocate(bI(par%ax%n,par%ay%n,par%az%n))
    allocate(bS(par%ax%n,par%ay%n,par%az%n))
    allocate(bV(par%ax%n,par%ay%n,par%az%n))
    allocate(bVr(par%nvrf,par%az%n))
    !allocate(dI(par%ax%n,par%ay%n,par%az%n))
    !allocate(dS(par%ax%n,par%ay%n,par%az%n))
    allocate(D(par%ax%n,par%ay%n,par%wblock%ns))
    allocate(bP(par%ax%n,par%ay%n,par%az%n,par%wblock%ns))
    allocate(dP_slice(par%ax%n,par%ay%n,nth))
    allocate(dU_slice(par%ax%n,par%ay%n,nth))
    allocate(kx2(par%ax%n),ky2(par%ay%n))
    allocate(tapx(par%ntpx),tapy(par%ntpy))

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

    allocate(cwrk1d(par%ax%n*5+par%ay%n*5,nth)) !!.. ifd or ffd
    allocate(cwrk3d(par%ax%n,par%ay%n,2,nth))   !!.. ssf-epi
    allocate(iwrk2d(par%ax%n,par%ay%n,nth))     !!.. ssf-epi
    allocate(rwrk1d(par%ax%n+par%ay%n,nth))
    stat = 0
  end function
  !!----------------------------------------------------------------------------------------
  integer function posmva_bgslow_init() result(stat)
    integer :: ix, iy

    call sep_read(bS(par%jx1:par%jx2,par%jy1:par%jy2,:), par%bS_)

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

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

    !!.. postack, so twice slowness
    bS = bS * 2.
    
    stat = 0
  end function
  !!-----------------------------------------------------------------------------------------
  integer function posmva_wfdexp_init() result(stat)
    integer :: iz

    bV = 1./bS 

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

    stat = wem_wfdexp_init(kx2, ky2, par%ax%n, par%ay%n, par%ax%d, par%ay%d, par%az%d, &
                           par%pdip, par%nvrf, par%oper, par%trick)
    stat = wem_tapers_init(tapx, tapy, par%ntpx, par%ntpy)
    stat = mva_phsscr_init(par%ax%d, par%ay%d, par%az%d, par%ax%n, par%ay%n)

    call k2_map(kx2, par%ax%n, par%ax%d)
    call k2_map(ky2, par%ay%n, par%ay%d)
    call wem_fftw_init(plan_2ddPF, plan_2ddPI, dP_slice, par%ax%n, par%ay%n, nth)

    stat = 0
  end function
  !!-----------------------------------------------------------------------------------------
  integer function posmva_bgwfld_oper() result(stat)
    type(axis) :: aws
    integer    :: ith, iwb, iws, iz
    real       :: ws
    ith = 1

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

    do iwb=1, par%wblock%nb
       aws%o = par%wblock%os(iwb); aws%d = par%wblock%ds; aws%n = par%wblock%blen(iwb)
       call sep_read(D(par%jx1:par%jx2,par%jy1:par%jy2,1:aws%n), par%D_)
       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/nwstep*nwstep==iws) write(0,*)"iws=", iws, ws, ith
          dP_slice(:,:,ith) = D(:,:,iws)
          do iz=1, par%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), par%ax%n, par%ay%n)
             bP(:,:,iz,iws) = dP_slice(:,:,ith)
             bI(:,:,iz) = bI(:,:,iz) + dP_slice(:,:,ith)/par%aw%n
          end do
       end do
!$OMP END PARALLEL DO
       call sep_write(bP(:,:,:,1:aws%n), par%bP_)
    end do
    call sep_write(bI(par%jx1:par%jx2,par%jy1:par%jy2,:), par%bI_)
    stat = 0
  end function
  !!-------------------------------------------------------------------------------------------
  integer function posmva_tomgph_oper(adj, add, dS, dI) result(stat)
    logical               :: adj, add
    real   , dimension(:) :: dS, dI

    if (.not. add) then
       if (adj) then
          dS = 0.
       else
          dI = 0.
       end if
    end if
    call posmva_wetomo(adj, dS, dI) 
    stat = 0
  end function
  !!--------------------------------------------------------------------------------------------
  subroutine posmva_wetomo(adj, dS, dI)
    logical                                     :: adj
    real, dimension(par%ax%n,par%ay%n,par%az%n) :: dS, dI 

    type(axis)                                  :: aws
    integer                                     :: ith, iwb, iws, iz, stat
    real                                        :: ws

    ith = 1

    call sseek(par%bP_, 0, 0)

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

       !write(0,*)"======================="
       !write(0,*)" back-projection... "
       !write(0,*)"======================="

       do iwb=1, par%wblock%nb
          aws%o = par%wblock%os(iwb); aws%d = par%wblock%ds; aws%n = par%wblock%blen(iwb)
          call sep_read(bP(:,:,:,1:aws%n), par%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/nwstep*nwstep==iws) write(0,*)"iws=", iws, ws, ith
             dP_slice(:,:,ith) = 0.
             dU_slice(:,:,ith) = 0.
             do iz=par%az%n, 1, -1

                dP_slice(:,:,ith) = dP_slice(:,:,ith) + dI(:,:,iz)

                stat = mva_phsscr_oper(dU_slice(:,:,ith), bP(:,:,iz,iws), bS(:,:,iz), ws)

                dS(:,:,iz) = dS(:,:,iz) + conjg(dU_slice(:,:,ith))*dP_slice(:,:,ith)/par%aw%n

                stat = wem_tapers_oper(dP_slice(:,:,ith), par%ax%n, par%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

       !write(0,*)"======================="
       !write(0,*)" forward-scattering... "
       !write(0,*)"======================="

       do iwb=1, par%wblock%nb
          aws%o = par%wblock%os(iwb); aws%d = par%wblock%ds; aws%n = par%wblock%blen(iwb)
          call sep_read(bP(:,:,:,1:aws%n), par%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/nwstep*nwstep==iws) write(0,*)"iws=", iws, ws, ith
             dU_slice(:,:,ith) = 0.
             dP_slice(:,:,ith) = 0.
             do iz=1, par%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), par%ax%n, par%ay%n)

                stat = mva_phsscr_oper(dU_slice(:,:,ith), bP(:,:,iz,iws), bS(:,:,iz), ws)

                dP_slice(:,:,ith) = dP_slice(:,:,ith) + dU_slice(:,:,ith)*dS(:,:,iz)

                dI(:,:,iz)  = dI(:,:,iz)  + dP_slice(:,:,ith)/par%aw%n

             end do
          end do
!$OMP END PARALLEL DO
       end do
    end if
  end subroutine


end module
