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


  integer                         , private :: indx, indy, indz, nw, nsh
  type(out_grid)                  , private :: ogrid
  type(pro_grid)                  , private :: pgrid 
  integer, dimension(:,:), pointer, private :: index_mxs, index_mxr
  integer, dimension(:,:), pointer, private :: index_mys, index_myr
  integer, dimension(:,:), pointer, private :: index_mzs, index_mzr
  integer, external               , private :: sreed_window, srite_window

  public                                    :: wem_imgcon_init
  public                                    :: wem_imgcon_oper
  public                                    :: wem_imgcon_clos
  private                                   :: imgcon_kernel
contains

  integer function wem_imgcon_init(ogrid_in, pgrid_in, nw_in, nsh_in,                                  &
                                   index_mxs_in, index_mxr_in, index_mys_in, index_myr_in) result(stat)
    type(out_grid)                  :: ogrid_in
    type(pro_grid)                  :: pgrid_in
    integer                         :: nw_in
    integer                         :: nsh_in
    integer, dimension(:,:), target :: index_mxs_in, index_mxr_in
    integer, dimension(:,:), target :: index_mys_in, index_myr_in

    integer                         :: imx, imy, imz
    integer                         :: ihx, ihy, ihz
    integer                         :: jhx, jhy, jhz
    real                            :: hx, hy, hz

    stat = 1

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

    ogrid = ogrid_in
    pgrid = pgrid_in

    nw  = nw_in
    nsh = nsh_in

    index_mxs => index_mxs_in
    index_mxr => index_mxr_in
    index_mys => index_mys_in
    index_myr => index_myr_in 

    do ihx=1, ogrid%ahx%n
       hx  = ogrid%ahx%o+(ihx-1)*ogrid%ahx%d
       jhx = nint(hx/pgrid%ax%d)
       do imx=1, ogrid%ax%n
          index_mxs(ihx,imx) = max(1, min(indx+imx-1-jhx,pgrid%ax%n))
          index_mxr(ihx,imx) = max(1, min(indx+imx-1+jhx,pgrid%ax%n))
       end do 
    end do

    do ihy=1, ogrid%ahy%n
       hy  = ogrid%ahy%o+(ihy-1)*ogrid%ahy%d
       jhy = nint(hy/pgrid%ay%d)
       do imy=1, ogrid%ay%n
          index_mys(ihy,imy) = max(1, min(indy+imy-1-jhy,pgrid%ay%n))
          index_myr(ihy,imy) = max(1, min(indy+imy-1+jhy,pgrid%ay%n))
       end do
    end do

    stat = 0

  end function

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

    indx = 0; indy = 0; indz = 0

    nullify(index_mxs, index_mxr)
    nullify(index_mys, index_myr)

    stat = 0
  end function

  integer function wem_imgcon_oper(adj, modl, wfds, wfdr, ws, pside) result(stat)
    logical                     :: adj
    real   , dimension(:,:,:,:) :: modl
    complex, dimension(:,:)     :: wfds !!.. wfds(nx, ny, nz)
    complex, dimension(:,:)     :: wfdr
    real                        :: ws
    character(len=3)            :: pside

    complex                     :: phase_shaping
    real                        :: PI

    PI = 3.14159265359

    stat = 1

    !!.. in case of 2D migration, 45 degree rotation i^(-1/2)
    phase_shaping = cmplx(cos(PI/4.), -sin(PI/4.))!/sqrt(ws)

    if (pgrid%ay%n>1 .and. pgrid%ax%n>1) then
       !!.. in case of 3D migration, 90 degree rotation i^(-1)
       phase_shaping = cmplx(0., - 1. * ws)
    end if

    call imgcon_kernel(adj, modl, wfds, wfdr, phase_shaping, pside)

    stat = 0
  end function

  !!.. imaging condition with all spatial lags
  subroutine imgcon_kernel(adj, modl, wfds, wfdr, phase_shaping, pside)
    logical                     :: adj
    real   , dimension(:,:,:,:) :: modl
    complex, dimension(:,:)     :: wfds !!.. wfds(nx, ny, nz)
    complex, dimension(:,:)     :: wfdr
    complex                     :: phase_shaping
    character(len=3)            :: pside 

    integer                     :: imx, imy, imz
    integer                     :: ihx, ihy, ihz


    if (adj.and.pside=="sou") then

       do imy=1, ogrid%ay%n
       do imx=1, ogrid%ax%n
       do ihy=1, ogrid%ahy%n
       do ihx=1, ogrid%ahx%n
          wfds(index_mxs(ihx,imx),index_mys(ihy,imy)) = wfds(index_mxs(ihx,imx),index_mys(ihy,imy))        &
                                                      + conjg(wfdr(index_mxr(ihx,imx),index_myr(ihy,imy))) &
                                                      * modl(ihx,ihy,imx,imy)*conjg(phase_shaping)
       end do
       end do
       end do
       end do 

    else if (adj.and.pside=="rec") then

       do imy=1, ogrid%ay%n
       do imx=1, ogrid%ax%n
       do ihy=1, ogrid%ahy%n
       do ihx=1, ogrid%ahx%n
          wfdr(index_mxr(ihx,imx),index_myr(ihy,imy)) = wfdr(index_mxr(ihx,imx),index_myr(ihy,imy))        &
                                                      + conjg(wfds(index_mxs(ihx,imx),index_mys(ihy,imy))) &
                                                      * modl(ihx,ihy,imx,imy)*conjg(phase_shaping)
       end do
       end do
       end do
       end do

    else

       do imy=1, ogrid%ay%n
       do imx=1, ogrid%ax%n
       do ihy=1, ogrid%ahy%n
       do ihx=1, ogrid%ahx%n
          modl(ihx,ihy,imx,imy) = modl(ihx,ihy,imx,imy)                                     &
                                + wfds(index_mxs(ihx,imx),index_mys(ihy,imy))               &
                                * wfdr(index_mxr(ihx,imx),index_myr(ihy,imy))*phase_shaping
       end do
       end do
       end do
       end do

    end if
    return
  end subroutine

  complex function conj(x)
    complex :: x
    real    :: xr, xi, xm
    xr = real(x)
    xi = aimag(x)
    xm = xr*xr + xi*xi
    conj = cmplx(xr*xr-xi*xi, -2*xr*xi)/xm
  end function

end module
