program main
  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_bornsc_mod
  implicit none

  type(axis)                                 :: ax, ay, az, aw, aws
  type(axis)                                 :: arx, ary, avx, avy
  type(block)                                :: wblock
  complex  , dimension(:,:,:,:), allocatable :: bP
  complex  , dimension(:,:,:)  , allocatable :: bP_slice
  complex  , dimension(:,:,:)  , allocatable :: dU_slice, dP_slice
  complex  , dimension(:,:,:)  , allocatable :: D
  real     , dimension(:,:,:)  , allocatable :: bI
  real     , dimension(:,:,:)  , allocatable :: dI
  real     , dimension(:,:,:)  , allocatable :: bS
  real     , dimension(:,:,:)  , allocatable :: dS
  real     , dimension(:,:,:)  , allocatable :: bV
  real     , dimension(:,:)    , allocatable :: bVr

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

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

  integer*8, dimension(:)      , allocatable :: plan_2dbPF, plan_2dbPI
  integer*8, dimension(:)      , allocatable :: plan_2ddPF, plan_2ddPI
  integer*8, dimension(:)      , allocatable :: plan_2ddUF, plan_2ddUI

  character(len=5)                           :: oper
  integer                                    :: nxpad_beg, nxpad_end
  integer                                    :: nypad_beg, nypad_end
  integer                                    :: pdip, nvrf, ntpx, ntpy
  real                                       :: trick, ws, PI
  real                                       :: rx_min, rx_max
  real                                       :: ry_min, ry_max
  real                                       :: px_min, px_max
  real                                       :: py_min, py_max
  real                                       :: f1, f2, f3, f4
  real                                       :: w1, w2, w3, w4

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

  logical                                    :: adj, radian
  integer                                    :: ix, iy, iz, iwb, iws, indw
  integer                                    :: nw1, nw2, nw3, nw4, iw
  integer                                    :: ith, nth, ompchunk, stat 
  integer                         , external :: omp_get_num_procs
  integer                         , external :: omp_get_num_threads
  integer                         , external :: omp_set_num_threads
  integer                         , external :: omp_get_thread_num

  call sep_init(SOURCE)
  call getch_add_string("head=/dev/null")
  call getch_add_string("noheader=y")
  call set_no_putch()
  call sep_begin_prog()

  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)

  call from_param("f1"       , f1       ,      5.)
  call from_param("f2"       , f2       ,     10.)
  call from_param("f3"       , f3       ,     35.)
  call from_param("f4"       , f4       ,     40.)

  !!.. 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)
  
  !!.. background wavefield
  call auxinout(bP_)

  if (adj) then !!.. back-projection
     call auxin (dI_)
     call auxout(dS_)
  else          !!.. forward scattering
     call auxout(dI_)
     call auxin (dS_)
  end if

  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

  call houaxis(bI_, arx, 1); call houaxis(bI_, ary, 2); call houaxis(bI_, az, 3)

  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_)

  if (adj) then
     call houaxis(dS_, arx, 1); call houaxis(dS_, ary, 2); call houaxis(dS_, az, 3)
  else
     call houaxis(dI_, arx, 1); call houaxis(dI_, ary, 2); call houaxis(dI_, az, 3)
  end if

  call base_blkcal(wblock, aw)

  allocate(bI(ax%n,ay%n,az%n))
  allocate(dI(ax%n,ay%n,az%n))
  allocate(bS(ax%n,ay%n,az%n))
  allocate(dS(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(bP_slice(ax%n,ay%n,nth))
  allocate(dU_slice(ax%n,ay%n,nth))
  allocate(dP_slice(ax%n,ay%n,nth))
  allocate(kx2(ax%n),ky2(ay%n))
  allocate(tapx(ntpx),tapy(ntpy))
  allocate(hamw(aw%n))

  allocate(plan_2dbPF(nth),plan_2dbPI(nth))
  allocate(plan_2ddUF(nth),plan_2ddUI(nth))
  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))     !!.. ifd or ffd

   
  if (adj) then
     call sep_read(bS(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,:), bS_)
  else
     call sep_read(dS(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,:), dS_)
     call sep_read(bS(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,:), bS_)
  end if

  !! 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 = mva_phsscr_init(ax%d, ay%d, az%d, ax%n, ay%n)

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

  w1 = 2*PI*f1; nw1 = 1.5+(w1-aw%o)/aw%d
  w2 = 2*PI*f2; nw2 = 1.5+(w2-aw%o)/aw%d 
  w3 = 2*PI*f3; nw3 = 1.5+(w3-aw%o)/aw%d
  w4 = 2*PI*f4; nw4 = 1.5+(w4-aw%o)/aw%d

write(0,*)nw1, nw2, nw3, nw4, aw%n


  do iw=1, aw%n
     if      (iw>=nw1 .and. iw<=nw2) then
        hamw(iw) = 0.54+0.46*cos(PI*(iw-nw1)/(nw2-nw1)-PI)
     else if (iw>=nw3 .and. iw<=nw4) then
        hamw(iw) = 0.54+0.46*cos(PI*(nw3-iw)/(nw4-nw3)   )
     else if (iw <nw1 .or.  iw> nw4) then
        hamw(iw) = 0.
     else 
        hamw(iw) = 1.
     end if
  end do


  call k2_map(kx2, ax%n, ax%d)
  call k2_map(ky2, ay%n, ay%d)
  call wem_fftw_init(plan_2dbPF, plan_2dbPI, bP_slice, ax%n, ay%n, nth)
  call wem_fftw_init(plan_2ddPF, plan_2ddPI, dP_slice, ax%n, ay%n, nth)
  call wem_fftw_init(plan_2ddUF, plan_2ddUI, dU_slice, ax%n, ay%n, nth) 

  !!.. compute the background wavefield and save it onto disk
  write(0,*)"Computing background wavefield..."
  bI = 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_)
    
     bP_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
        indw = 1.5 + (ws-aw%o)/aw%d
    
        if (iws/10*10==iws) write(0,*)"iws=", iws, ws, ith

        bP_slice(:,:,ith) = D(:,:,iws)
        do iz=1, az%n
           stat = wem_wfdexp_oper(bP_slice(:,:,ith), bV(:,:,iz), bVr(:,iz), ws, -1,                          &
                  iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dbPF(ith), plan_2dbPI(ith))
           stat = wem_tapers_oper(bP_slice(:,:,ith), ax%n, ay%n)

           bP(:,:,iz,iws) = bP_slice(:,:,ith)
           bI(:,:,iz) = bI(:,:,iz) + bP_slice(:,:,ith)/aw%n
        end do
     end do
!$OMP END PARALLEL DO 
     call sep_write(bP(:,:,:,1:aws%n), bP_)
     
  end do

  call sep_write(bI(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,:), bI_)

  call sseek(bP_, 0, 0)


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

     dI = 0. 
     call sep_read(dI(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,:), dI_)
!     iz = 1.5 + (600-ax%o)/ax%d
!     dI(1:iz-3,:,:) = 0.
!     dI(iz+3: ,:,:) = 0.
     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
           indw = 1.5 + (ws-aw%o)/aw%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
              bP_slice(:,:,ith) = bP(:,:,iz,iws)

              stat = mva_phsscr_oper(dU_slice(:,:,ith), bP_slice(:,:,ith), 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
     call sep_write(dS(nxpad_beg+1:nxpad_beg+arx%n,nypad_beg+1:nypad_beg+ary%n,:), dS_)


  else           !!.. dI = F  dS
     dI = 0.
     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
           indw = 1.5 + (ws-aw%o)/aw%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
              bP_slice(:,:,ith) = bP(:,:,iz,iws)

              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)

              stat = mva_phsscr_oper(dU_slice(:,:,ith), bP_slice(:,:,ith), bS(:,:,iz), ws)
              dU_slice(:,:,ith) = dU_slice(:,:,ith)*dS(:,:,iz)

              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

  stat = wem_velref_clos()
  stat = wem_wfdexp_clos()
  stat = wem_tapers_clos()

  call wem_fftw_clos(plan_2dbPF, plan_2dbPI, nth)
  call wem_fftw_clos(plan_2ddPF, plan_2ddPI, nth)
  call wem_fftw_clos(plan_2ddUF, plan_2ddUI, nth)
  deallocate(D)
  deallocate(bI, dI)
  deallocate(bS, dS)
  deallocate(bV, bVr)
  deallocate(bP, bP_slice)
  deallocate(dP_slice, dU_slice)
  deallocate(tapx, tapy, hamw)
  deallocate(kx2, ky2)
  deallocate(plan_2dbPF, plan_2dbPI)
  deallocate(plan_2ddPF, plan_2ddPI)
  deallocate(plan_2ddUF, plan_2ddUI)
end program
