!!---------------------------------------------------------------------
!!.. Wave-equation shot-profile migration, driver's module
!!.. OpenMP parallelized
!!---------------------------------------------------------------------
!!.. 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_mod
!  use  wem_glodef_mod !!.. global definition
  use  wem_glocal_mod !!.. global calculation
  use  wem_fileio_mod !!.. file I/O
  use  wem_velref_mod !!.. reference velocity
  use  wem_wfdexp_mod !!.. wavefield extrapolation
  use  wem_imgcon_mod !!.. imaging condition 
  use  wem_fftwv3_mod !!.. FFTW3
  use  wem_hesapp_mod !!.. Hessian
  use  wem_tapers_mod !!.. taper
  use  base_blkcal_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(hes_grid)                             , private :: hgrid

  type(block)                                , private :: wblock

  real     , dimension(:,:,:)   , allocatable, private :: vmod     !!.. velocity model
  real     , dimension(:,:)     , allocatable, private :: vref     !!.. reference velocity

  real     , dimension(:,:,:)   , allocatable, private :: modl     !!.. output image
  real     , dimension(:,:,:)   , allocatable, private :: hess     !!.. Hessian

  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 :: iwrk2d
  real     , dimension(:,:)     , allocatable, private :: rwrk1d
  complex  , dimension(:,:)     , allocatable, private :: cwrk1d
  complex  , dimension(:,:,:,:) , allocatable, private :: cwrk3d

  character(len=256)                         , private :: csou_
  character(len=256)                         , private :: crec_
  character(len=256)                         , private :: vmod_
  character(len=256)                         , private :: imag_
  character(len=256)                         , private :: hess_
  character(len=256)                         , private :: hsou_
contains

  integer function wem_shpmig_init() result(stat)
    stat = 1
    call wem_param_init(dgrid, pgrid, igrid, hgrid, wblock, crec_, csou_, vmod_, imag_, hess_, hsou_)
    call wem_histo_data(dgrid, vgrid, igrid, pgrid, crec_, csou_, vmod_, hsou_)
    call wem_histo_imag(igrid, imag_)
    if (hgrid%wnthes) call wem_histo_hess(hgrid, igrid, hess_)
    stat = 0
  end function

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

  integer function wem_shpmig_oper() result(stat)
    type(axis)        :: aws
    integer           :: iws, iwb, izs, ish
    real              :: ws, PI
    integer           :: ogrid_nx_bak
    integer           :: ogrid_ny_bak
    integer           :: ogrid_nz_bak
    integer           :: ith, nth, ompchunk
    integer, external :: omp_get_num_procs   
    integer, external :: omp_get_num_threads
    integer, external :: omp_set_num_threads
    integer, external :: omp_get_thread_num

    PI = 3.14159265359

    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

    !!.. 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,*)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 wem_glocal(pgrid, ogrid, vgrid, dgrid, igrid, hgrid, 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))                       !!.. source array
    allocate(crec(dgrid%ntrg,wblock%ns))                       !!.. receiver array
    allocate(wfds(pgrid%ax%n,pgrid%ay%n,pgrid%az%n,wblock%ns)) !!.. source wavefield
    allocate(wfdr(pgrid%ax%n,pgrid%ay%n,pgrid%az%n,wblock%ns)) !!.. receiver wavefield
    allocate(vmod(pgrid%ax%n,pgrid%ay%n,pgrid%az%n))           !!.. velocity model
    allocate(vref(pgrid%nvrf,pgrid%az%n))                      !!.. reference velocity
    if (igrid%wntimg) then
       allocate(modl(ogrid%ax%n,ogrid%ay%n,ogrid%az%n))
       call wem_fio_imag_init(imag_, igrid)
    end if
    if (hgrid%wnthes) then
       allocate(hess(ogrid%ax%n,ogrid%ay%n,ogrid%az%n))
       call wem_fio_hess_init(hess_, hgrid)
    end if
    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(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

    stat = wem_tapers_init(tapx, tapy, pgrid%ntpx, pgrid%ntpy)
    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)


    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
          if (igrid%wntimg) then
             deallocate(modl); allocate(modl(ogrid%ax%n,ogrid%ay%n,ogrid%az%n)); modl = 0.
          end if
          if (hgrid%wnthes) then
             deallocate(hess); allocate(hess(ogrid%ax%n,ogrid%ay%n,ogrid%az%n)); hess = 0.
          end if
       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, vmod_)
       do izs=1, pgrid%az%n
          stat = wem_velref_oper(vmod(:,:,izs), vref(:,izs))
       end do


       if (igrid%wntimg) stat = wem_imgcon_init(igrid, ogrid, pgrid, imag_) 
       if (hgrid%wnthes) stat = wem_hesapp_init(hgrid, ogrid, pgrid, hess_)
       do iwb=1, wblock%nb
          aws%o = wblock%os(iwb); aws%d = wblock%ds; aws%n = wblock%blen(iwb)
write(0,*)"aws%o=",aws%o,"aws%n=",aws%n
          call wem_fio_csou(csou, dgrid, wblock%indx(iwb), aws%n, ish, csou_, "reed")
          call wem_fio_crec(crec, dgrid, wblock%indx(iwb), aws%n, ish, crec_, "reed")
          call wem_wfs_init(wfds, csou, dgrid, pgrid)
          call wem_wfr_init(wfdr, crec, dgrid, pgrid)

!$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/10*10==iws) write(0,*)"iws=", iws, ws, ith

             wfds_slice(:,:,ith) = wfds(:,:,1,iws)
             wfdr_slice(:,:,ith) = wfdr(:,:,1,iws)
             do izs=2, pgrid%nzstep !pgrid%az%n
                stat = wem_wfdexp_oper(wfds_slice(:,:,ith), vmod(:,:,izs-1), vref(:,izs-1), 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-1), vref(:,izs-1), 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(:,:,izs,iws) = wfds_slice(:,:,ith)
                wfdr(:,:,izs,iws) = wfdr_slice(:,:,ith)
             end do
!$OMP CRITICAL
             if (igrid%wntimg) stat = wem_imgcon_oper(modl, wfds(:,:,:,iws), wfdr(:,:,:,iws), ws)
             if (hgrid%wnthes) stat = wem_hesapp_oper(hess, wfds(:,:,:,iws), wfdr(:,:,:,iws), ws) 
!$OMP END CRITICAL

          end do
!$OMP END PARALLEL DO
       end do

       stat = wem_velref_clos()
       stat = wem_wfdexp_clos()
       if (igrid%wntimg) stat = wem_imgcon_clos()
       if (hgrid%wnthes) stat = wem_hesapp_clos()
    end do

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

    stat = wem_tapers_clos()

    if (igrid%wntimg) deallocate(modl)
    if (hgrid%wnthes) deallocate(hess)

    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)

    stat = 0
  end function

end module

