!!---------------------------------------------------------------------
!!.. Wave-equation shot-profile migration, driver's module
!!.. OpenMP parallelized
!!.. Version 4.0
!!.. This version includes the forward and adjoint operators
!!.. This version also parallelizes the imaging condition
!!---------------------------------------------------------------------
!!.. Author: Yaxun Tang, tang@sep.stanford.edu
!!---------------------------------------------------------------------
!!.. Note: 1. To make full usage of the threads, 
!!..          it is better to keep nws >= nth 
!!..       2. For threads safty purpose, the memory allocation is
!!..          considerably increased compared to serial implementation.
!!..          Mem of parallel ~ Mem of serial * nth
!!---------------------------------------------------------------------
module wem_shpmig_complex_mod
  use  base_blkcal_mod
  use  wem_struct_mod
  use  wem_histos_complex_mod !!.. history files
  use  wem_glocal_mod         !!.. global calculation
  use  wem_fileio_complex_mod !!.. file I/O
  use  wem_velref_mod         !!.. reference velocity
  use  wem_wfdexp_mod         !!.. wavefield extrapolation
  use  wem_imgcon_complex_mod !!.. imaging condition 
  use  wem_fftwv3_mod         !!.. FFTW3
  use  wem_tapers_mod         !!.. taper
  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(hes_grid)                             , private :: hgrid

  type(block)                                , private :: wblock

  type(wem_ftag)                             , private :: ftag_

  real     , dimension(:,:,:)    , allocatable, private :: vmod     !!.. velocity model
  real     , dimension(:,:)      , allocatable, private :: vref     !!.. reference velocity
  complex  , dimension(:,:,:,:,:), allocatable, private :: modl     !!.. output image

  complex  , dimension(:,:)      , allocatable, private :: crec
  complex  , dimension(:,:)      , allocatable, private :: csou
  complex  , dimension(:,:,:,:)  , allocatable, private :: wfds
  complex  , dimension(:,:,:,:)  , allocatable, private :: wfdr

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

  complex  , dimension(:,:,:)    , allocatable, private :: wfds_slice
  complex  , dimension(:,:,:)    , allocatable, private :: wfdr_slice

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

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

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

  integer                                     , private :: nzrept

  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 wem_shpmig_init(adj, dgrid_out, igrid_out, ftag_out_) result(stat)
    logical           :: adj
    type(dat_grid)    :: dgrid_out
    type(img_grid)    :: igrid_out
    type(wem_ftag)    :: ftag_out_

    stat = 1

    ftag_%csou_ = "csou"
    ftag_%crec_ = "crec"
    ftag_%vmod_ = "vmod"
    ftag_%imag_ = "imag"
    ftag_%hsou_ = "hsou"

!   read parameters
!   ---------------
    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("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("nzrept", nzrept      , 50     )

!   initialize file headers
!   -----------------------
    call auxinout(ftag_%csou_)
    call auxinout(ftag_%crec_)
    call auxinout(ftag_%vmod_)
    call auxinout(ftag_%imag_)
    if (dgrid%hdrsou) call auxinout(ftag_%hsou_)

    call wem_histo_data(adj, dgrid, vgrid, igrid, pgrid, ftag_)
    call wem_histo_imag(adj, igrid, vgrid, ftag_)

!   setup OpenMP stuff
!   ------------------
    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_)
    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

!   calculate the global parameters for the first shot
!   --------------------------------------------------
    call wem_glocal(pgrid, ogrid, vgrid, dgrid, igrid, hgrid, 1)

!   calculate the frequency blocks
!   ------------------------------
    call base_blkcal(wblock, dgrid%aw)

    allocate(csou(dgrid%ntrs,wblock%ns))                              !!.. source array
    allocate(crec(dgrid%ntrg,wblock%ns))                              !!.. receiver array
    allocate(wfds(pgrid%ax%n,pgrid%ay%n,wblock%ns,pgrid%az%n+1))      !!.. source wavefield
    allocate(wfdr(pgrid%ax%n,pgrid%ay%n,wblock%ns,pgrid%az%n+1))      !!.. receiver wavefield
    allocate(vmod(pgrid%ax%n,pgrid%ay%n,pgrid%az%n))                  !!.. velocity model
    allocate(vref(pgrid%nvrf,pgrid%az%n))                             !!.. reference velocity
    allocate(modl(ogrid%ahx%n,ogrid%ahy%n,ogrid%ax%n,ogrid%ay%n,nth))
    allocate(wfds_slice(pgrid%ax%n,pgrid%ay%n,nth))
    allocate(wfdr_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(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))     !!.. ifd or ffd

    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, wfds_slice, pgrid%ax%n, pgrid%ay%n, nth)
    call wem_fftw_init(plan_2drF, plan_2drI, wfdr_slice, pgrid%ax%n, pgrid%ay%n, nth)

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

    dgrid_out = dgrid
    igrid_out = igrid
    ftag_out_ = ftag_

    stat = 0
  end function

  integer function wem_shpmig_clos() result(stat)
    stat = 0
  end function

  integer function wem_shpmig_oper(adj, add, imag_, crec_) result(stat)
    logical           :: adj
    logical           :: add
    character(len=*)  :: imag_
    character(len=*)  :: crec_ 

    type(axis)        :: aws
    integer           :: iws, iwb, izs, ish, ith
    real              :: ws, PI
    integer           :: ogrid_nx_bak
    integer           :: ogrid_ny_bak
    integer           :: ogrid_nz_bak


    if (.not.add) then
       if (adj) then
          call wem_fio_imag_init(igrid, imag_)
       else
          call wem_fio_crec_init(dgrid, crec_)
       end if 
    end if

!   shot-profile migration 
!   ----------------------
    if (adj) then
       do ish=1, dgrid%nshg
          ogrid_nx_bak = ogrid%ax%n; ogrid_ny_bak = ogrid%ay%n; ogrid_nz_bak = ogrid%az%n

!         calculate the global parameters
!         -------------------------------
          call wem_glocal(pgrid, ogrid, vgrid, dgrid, igrid, hgrid, 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(modl)
             deallocate(index_mxs,index_mxr)
             deallocate(index_mys,index_myr)
             allocate(modl(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,*)"pgrid =", pgrid%ax%n,pgrid%ay%n,pgrid%az%n
          write(0,*)"ogrid%ax =", ogrid%ax%n,ogrid%ax%o,ogrid%ax%d
          write(0,*)"ogrid%ay =", ogrid%ay%n,ogrid%ay%o,ogrid%ay%d
          write(0,*)"ogrid%az =", ogrid%az%n,ogrid%az%o,ogrid%az%d
          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 wem_fio_velo(vmod, vgrid, pgrid, 1, pgrid%az%n, ftag_%vmod_, "reed")
          do izs=1, pgrid%az%n
             stat = wem_velref_oper(vmod(:,:,izs), vref(:,izs))
          end do
                       
          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)
             wfds = 0; wfdr = 0.
             call wem_fio_csou(csou, dgrid, wblock%indx(iwb), aws%n, ish, ftag_%csou_, "reed")
             call wem_fio_crec(crec, dgrid, wblock%indx(iwb), aws%n, ish,       crec_, "reed")
             call wem_wfs_init(.true., .false., wfds(:,:,:,1), csou, dgrid, pgrid)
             call wem_wfr_init(.true., .false., wfdr(:,:,:,1), crec, dgrid, pgrid)

             do izs=1, pgrid%az%n
                modl = 0.
                call wem_fio_imag(modl(:,:,:,:,1), igrid, ogrid, izs, imag_, "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
                   wfds_slice(:,:,ith) = wfds(:,:,iws,izs)
                   wfdr_slice(:,:,ith) = wfdr(:,:,iws,izs)
                   stat = wem_imgcon_oper(.false., modl(:,:,:,:,ith), wfds_slice(:,:,ith), wfdr_slice(:,:,ith), ws, "all")
                   stat = wem_wfdexp_oper(wfds_slice(:,:,ith), vmod(:,:,izs), vref(:,izs), ws, -1,                         &
                          iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dsF(ith), plan_2dsI(ith))
                   stat = wem_wfdexp_oper(wfdr_slice(:,:,ith), vmod(:,:,izs), vref(:,izs), ws, -1,                         &
                          iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2drF(ith), plan_2drI(ith))
                   stat = wem_tapers_oper(wfds_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                   stat = wem_tapers_oper(wfdr_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                   wfds(:,:,iws,izs+1) = wfds_slice(:,:,ith)
                   wfdr(:,:,iws,izs+1) = wfdr_slice(:,:,ith)
                end do
                !$OMP END PARALLEL DO
                call wem_fio_imag(sum(modl,5), igrid, ogrid, izs, imag_, "rite")  
             end do
          end do
          stat = wem_velref_clos()
          stat = wem_wfdexp_clos()
          stat = wem_imgcon_clos()
       end do

    else

!      shot-profile demigration
!      ------------------------

       do ish=1, dgrid%nshg
          ogrid_nx_bak = ogrid%ax%n; ogrid_ny_bak = ogrid%ay%n; ogrid_nz_bak = ogrid%az%n

          call wem_glocal(pgrid, ogrid, vgrid, dgrid, igrid, hgrid, 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
             deallocate(modl)
             deallocate(index_mxs,index_mxr)
             deallocate(index_mys,index_myr)
             allocate(modl(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,*)"pgrid =", pgrid%ax%n,pgrid%ay%n,pgrid%az%n
          write(0,*)"ogrid%ax =", ogrid%ax%n,ogrid%ax%o,ogrid%ax%d
          write(0,*)"ogrid%ay =", ogrid%ay%n,ogrid%ay%o,ogrid%ay%d
          write(0,*)"ogrid%az =", ogrid%az%n,ogrid%az%o,ogrid%az%d
          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 wem_fio_velo(vmod, vgrid, pgrid, 1, pgrid%az%n, ftag_%vmod_, "reed")
          do izs=1, pgrid%az%n
             stat = wem_velref_oper(vmod(:,:,izs), vref(:,izs))
          end do
          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)
             wfds = 0.; wfdr = 0.
             call wem_fio_csou(csou, dgrid, wblock%indx(iwb), aws%n, ish, ftag_%csou_, "reed")
             call wem_wfs_init(.true., .false., wfds(:,:,:,1), csou, dgrid, pgrid)

             do izs=1, pgrid%az%n-1
                !$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
                   wfds_slice(:,:,ith) = wfds(:,:,iws,izs)
                   stat = wem_wfdexp_oper(wfds_slice(:,:,ith), vmod(:,:,izs), vref(:,izs), ws, -1,                         &
                          iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2dsF(ith), plan_2dsI(ith))
                   stat = wem_tapers_oper(wfds_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                   wfds(:,:,iws,izs+1) = wfds_slice(:,:,ith)
                end do
                !$OMP END PARALLEL DO 
             end do

             do izs=pgrid%az%n, 1, -1
                call wem_fio_imag(modl(:,:,:,:,1), igrid, ogrid, izs, imag_, "reed")
                do ith=1, nth
                   modl(:,:,:,:,ith)  = modl(:,:,:,:,1)
                end do
                !$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
                   wfds_slice(:,:,ith) = wfds(:,:,iws,izs  )
                   wfdr_slice(:,:,ith) = wfdr(:,:,iws,izs+1)
                   stat = wem_wfdexp_oper(wfdr_slice(:,:,ith), vmod(:,:,izs), vref(:,izs), ws, +1,                         &   
                          iwrk2d(:,:,ith), rwrk1d(:,ith), cwrk1d(:,ith), cwrk3d(:,:,:,ith), plan_2drF(ith), plan_2drI(ith))
                   stat = wem_tapers_oper(wfdr_slice(:,:,ith), pgrid%ax%n, pgrid%ay%n)
                   stat = wem_imgcon_oper(.true., modl(:,:,:,:,ith), wfds_slice(:,:,ith), wfdr_slice(:,:,ith), ws, "rec")
                   wfdr(:,:,iws,izs) = wfdr_slice(:,:,ith)
                end do
                !$OMP END PARALLEL DO
             end do

             if (add) call wem_fio_crec(crec, dgrid, wblock%indx(iwb), aws%n, ish, crec_, "reed")
             call wem_wfr_init(.false., add, wfdr(:,:,:,1), crec, dgrid, pgrid)
             call wem_fio_crec(crec, dgrid, wblock%indx(iwb), aws%n, ish, crec_, "rite")
          end do
          stat = wem_velref_clos()
          stat = wem_wfdexp_clos()
          stat = wem_imgcon_clos()
       end do
    end if

    call wem_fftw_clos(plan_2dsF, plan_2dsI, nth)
    call wem_fftw_clos(plan_2drF, plan_2drI, nth)

    stat = wem_tapers_clos()

    deallocate(modl)
    deallocate(csou, crec)
    deallocate(wfds, wfdr)
    deallocate(vmod, vref)
    deallocate(wfds_slice)
    deallocate(wfdr_slice)
    deallocate(tapx, tapy)
    deallocate(kx2 , ky2 )
    deallocate(plan_2dsF, plan_2dsI)
    deallocate(plan_2drF, plan_2drI)
    deallocate(index_mxs, index_mxr)
    deallocate(index_mys, index_myr)
    stat = 0
  end function

end module

