!!.. This module contains various kinds of imaging conditions
!!----------------------------------------------------------------
!!.. Author: Yaxun Tang, tang@sep.stanford.edu
!!----------------------------------------------------------------
module wem_imgcon_mod
  use  base_struct_mod
  implicit none


  integer                       , private :: indx, indy, indz
  type(img_grid)                , private :: igrid
  type(out_grid)                , private :: ogrid
  type(pro_grid)                , private :: pgrid 
  character(len=256)            , private :: filetag
  integer, external             , private :: sreed_window, srite_window

contains

  integer function wem_imgcon_init(igrid_in, ogrid_in, pgrid_in, filetag_in) result(stat)
    type(img_grid)        :: igrid_in
    type(out_grid)        :: ogrid_in
    type(pro_grid)        :: pgrid_in
    character(len=*)      :: filetag_in

    stat = 1

    indx = ogrid_in%index_outpro_xmin
    indy = ogrid_in%index_outpro_ymin
    indz = ogrid_in%index_outpro_zmin

    igrid = igrid_in
    ogrid = ogrid_in
    pgrid = pgrid_in

    filetag = filetag_in

    stat = 0

  end function

  integer function wem_imgcon_clos() result(stat)
    stat = 1

    indx = 0; indy = 0; indz = 0

    stat = 0
  end function

  integer function wem_imgcon_oper(modl, wfds, wfdr, ws) result(stat)
    real   , dimension(:,:,:) :: modl
    complex, dimension(:,:,:) :: wfds !!.. wfds(nx, ny, nz)
    complex, dimension(:,:,:) :: wfdr
    real                      :: ws

    complex                   :: phase_shaping
    real                      :: PI

    PI = 3.14159265359

    stat = 1

    if (pgrid%ay%n==1 .or. pgrid%ax%n==1) then
       !!.. in case of 2D migration, 45 degree rotation i^(-1/2)
       phase_shaping = cmplx(cos(PI/4.), -sin(PI/4.)) / sqrt(ws)
    else
       !!.. in case of 3D migration, 90 degree rotation i^(-1)
       phase_shaping = cmplx(0., - 1. * ws)
    end if

    call imgcon_kernel(modl, wfds, wfdr, phase_shaping)

    stat = 0
  end function


  !!.. imaging condition with all spatial lags
  subroutine imgcon_kernel(modl, wfds, wfdr, phase_shaping)
    real   , dimension(:,:,:) :: modl
    complex, dimension(:,:,:) :: wfds !!.. wfds(nx, ny, nz)
    complex, dimension(:,:,:) :: wfdr
    complex                   :: phase_shaping

    integer                   :: imx, imy, imz
    integer                   :: ihx, ihy, ihz
    integer                   :: jhx, jhy, jhz
    integer                   :: index_mxs, index_mys, index_mzs
    integer                   :: index_mxr, index_myr, index_mzr
    real                      :: hx, hy, hz
    complex                   :: wfld_sou, wfld_rec


    do ihz=1, igrid%ahz%n; hz = igrid%ahz%o+(ihz-1)*igrid%ahz%d; jhz = nint(hz/pgrid%az%d)
    do ihy=1, igrid%ahy%n; hy = igrid%ahy%o+(ihy-1)*igrid%ahy%d; jhy = nint(hy/pgrid%ay%d)
    do ihx=1, igrid%ahx%n; hx = igrid%ahx%o+(ihx-1)*igrid%ahx%d; jhx = nint(hx/pgrid%ax%d)

       call imgcon_fileio(modl, igrid, ogrid, ihx, ihy, ihz, filetag, "reed")

       do imz=1, ogrid%az%n; index_mzs = max(1, min(indz+imz-1-jhz,pgrid%az%n)); index_mzr = max(1, min(indz+imz-1+jhz,pgrid%az%n))
       do imy=1, ogrid%ay%n; index_mys = max(1, min(indy+imy-1-jhy,pgrid%ay%n)); index_myr = max(1, min(indy+imy-1+jhy,pgrid%ay%n))
       do imx=1, ogrid%ax%n; index_mxs = max(1, min(indx+imx-1-jhx,pgrid%ax%n)); index_mxr = max(1, min(indx+imx-1+jhx,pgrid%ax%n))

          wfld_sou = conjg(wfds(index_mxs,index_mys,index_mzs))
          wfld_rec =       wfdr(index_mxr,index_myr,index_mzr)
          modl(imx,imy,imz) = modl(imx,imy,imz) + real(wfld_sou*wfld_rec*phase_shaping)

       end do
       end do   
       end do

       call imgcon_fileio(modl, igrid, ogrid, ihx, ihy, ihz, filetag, "rite")

    end do
    end do
    end do
    return
  end subroutine

  subroutine imgcon_fileio(modl, igrid, ogrid, index_hx, index_hy, index_hz, filetag, flag)
    real, dimension(:,:,:) :: modl
    type(img_grid)         :: igrid
    type(out_grid)         :: ogrid
    integer                :: index_hx, index_hy, index_hz
    character(len=*)       :: filetag
    character(len=4)       :: flag

    integer, dimension(6)  :: n_cube, j_wind, n_wind, f_wind
    integer                :: n_dim_cube, err

    n_dim_cube = 6
    n_cube = (/igrid%ax%n,igrid%ay%n,igrid%az%n,igrid%ahx%n,igrid%ahy%n,igrid%ahz%n/)
    j_wind = (/1,1,1,1,1,1/)
    n_wind = (/ogrid%ax%n,ogrid%ay%n,ogrid%az%n,1,1,1/)
    f_wind = (/ogrid%index_outimg_xmin-1,ogrid%index_outimg_ymin-1,ogrid%index_outimg_zmin-1,index_hx-1,index_hy-1,index_hz-1/)

    if (flag == "rite") then
       err = srite_window(filetag,n_dim_cube,n_cube,n_wind,f_wind,j_wind,4,modl)
    else if (flag == "reed") then
       modl = 0.
       err = sreed_window(filetag,n_dim_cube,n_cube,n_wind,f_wind,j_wind,4,modl)
    end if

    return
  end subroutine


end module
