!!---------------------------------------------------------------------
!!.. 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 grn_oneway_mod
!  use  wem_glodef_mod !!.. global definition
  use  grn_glocal_mod !!.. global calculation
  use  grn_fileio_mod !!.. file I/O
  use  wem_velref_mod !!.. reference velocity
  use  wem_wfdexp_mod !!.. wavefield extrapolation
  use  wem_fftwv3_mod !!.. FFTW3
  use  wem_tapers_mod !!.. taper
  use  base_velgrd_mod
  use  base_blkcal_mod
  implicit none

  type(vel_grid)                             , private :: vgrid    !!.. velocity grid
  type(pro_grid)                             , private :: pgrid    !!.. propagation grid
  type(gre_grid)                             , private :: ggrid    !!.. image grid

  type(block)                                , private :: wblock

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

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

  complex  , dimension(:)       , allocatable, private :: csou
  complex  , dimension(:,:,:,:) , allocatable, private :: wfds
  complex  , dimension(:,:,:)   , allocatable, private :: wfds_slice

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

  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 :: vmod_
  character(len=256)                         , private :: wfds_

contains

  integer function grn_oneway_init() result(stat)
    stat = 1
    call grn_param_init(ggrid, pgrid, wblock, csou_, vmod_, wfds_)
    call grn_histo_data(ggrid, pgrid, vgrid, csou_, vmod_)
    call grn_histo_wfds(ggrid, wfds_)
    stat = 0
  end function

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

  integer function grn_oneway_oper() result(stat)
    type(axis)        :: aws
    integer           :: iws, iwb, izs, ish
    integer           :: ix_surf, iy_surf
    real              :: ws, x_surf, y_surf
    integer           :: ith, nth, ompchunk
    integer, external :: omp_get_num_threads
    integer, external :: omp_set_num_threads
    integer, external :: omp_get_thread_num

    ompchunk = 1
    nth      = 1
    ith      = 1

!$OMP PARALLEL
    if (omp_get_thread_num() == 0) then
       nth = omp_get_num_threads()
       write(0,*)"nthread =", nth
    end if
!$OMP END PARALLEL

    call grn_progrd_gen3d(pgrid, vgrid, ggrid)
    call base_velgrd_gen3d(vgrid, pgrid)
    call base_blkcal(wblock, ggrid%aw)

write(0,*)"pgrid_x:", pgrid%ax%n, pgrid%ax%o, pgrid%ax%d
write(0,*)"pgrid_y:", pgrid%ay%n, pgrid%ay%o, pgrid%ay%d
write(0,*)"pgrid_z:", pgrid%az%n, pgrid%az%o, pgrid%az%d 

    write(0,*)"wblock%nb", wblock%nb

    allocate(csou(wblock%ns))                                  !!.. source array
    allocate(wfds(pgrid%ax%n,pgrid%ay%n,pgrid%az%n,wblock%ns)) !!.. source wavefield
    allocate(vmod(pgrid%ax%n,pgrid%ay%n,pgrid%az%n))           !!.. velocity model
    allocate(vref(pgrid%nvrf,pgrid%az%n))                      !!.. reference velocity
    allocate(wfds_slice(pgrid%ax%n,pgrid%ay%n,nth))            !!.. a slice of wavefield
    allocate(kx2(pgrid%ax%n),ky2(pgrid%ay%n))                  !!.. wavenumber arrays
    allocate(tapx(pgrid%ntpx),tapy(pgrid%ntpy))                !!.. tapers

    allocate(plan_2dsF(nth),plan_2dsI(nth))                    !!.. FFTW plans

    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 grn_fio_wfds_init(ggrid, wfds_)
    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 grn_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

    write(0,*)"initialization finished..."

    do iy_surf=1, ggrid%ay_surf%n; y_surf=ggrid%ay_surf%o+(iy_surf-1)*ggrid%ay_surf%d; ggrid%index_surpro_sy=1.5+(y_surf-pgrid%ay%o)/pgrid%ay%d
    do ix_surf=1, ggrid%ax_surf%n; x_surf=ggrid%ax_surf%o+(ix_surf-1)*ggrid%ax_surf%d; ggrid%index_surpro_sx=1.5+(x_surf-pgrid%ax%o)/pgrid%ax%d
       write(0,*)"--------------------------------------------------"
       write(0,*)"ix_surf=", ix_surf, "      iy_surf", iy_surf
       write(0,*)"--------------------------------------------------"
       do iwb=1, wblock%nb
          aws%o = wblock%os(iwb); aws%d = wblock%ds; aws%n = wblock%blen(iwb)
          call grn_fio_csou(csou, ggrid, wblock%indx(iwb), aws%n, csou_)
          call grn_wfs_init(wfds, csou, ggrid, 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)
             do izs=1, pgrid%az%n
                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(:,:,izs,iws) = wfds_slice(:,:,ith)
             end do
          end do
!$OMP END PARALLEL DO
          call grn_fio_wfds(wfds, ggrid, wblock%indx(iwb), aws%n, ix_surf, iy_surf, wfds_, "rite")
       end do
    end do
    end do

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

    call wem_fftw_clos(plan_2dsF, plan_2dsI, nth)

    deallocate(csou, wfds)
    deallocate(vmod, vref)
    deallocate(wfds_slice)
    deallocate(tapx, tapy)
    deallocate( kx2,  ky2)
    deallocate(plan_2dsF, plan_2dsI)

    stat = 0
  end function

end module

