module weiop_mwc
  !! Mixed-domain Wavefield-Continuation
  !! fk-fx migration
  !! receive wfld at (izs,iws,ith)
  !! return  wfld downward/upward continued
  !! needs FKop, FXop

  use wei_util
  use wei_pro
  use wei_S
  use wei_fft
  use wei_comp
  use wei_dump
  !---------------
  implicit none

  !! temporary interpolation arrays
  complex, dimension(:,:,:,:,:), pointer, private :: tmp1,tmp2
  real,    dimension(:,:,:,:,:), pointer, private :: imap

contains
  !----------------------------------------------------------------
  !----------------------------------------------------------------
  integer function weiop_mwc0_init() result(st)
    integer :: er
    character(len=128), parameter :: unit='weiop_mwc0_init'
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif
    !! do nothing
#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    st=weiOK
  end function weiop_mwc0_init
  !----------------------------------------------------------------
  !! 1 reference slowness
  integer function weiop_mwc0(wfld,iws,izs,ith,FKop,FXop) result(st)
    complex, dimension(:,:,:,:,:), pointer :: wfld
    integer, intent(in)       :: iws,izs,ith
    integer                   :: ifk
    character(len=128), parameter :: unit='weiop_mwc0'
    interface !!----------------------------------------------
       integer function FKop(iws,izs,ifk,ith,wfld) result(st)
         integer, intent(in) :: iws,izs,ifk,ith
         complex, dimension(:,:,:,:,:), pointer :: wfld
       end function FKop
       integer function FXop(iws,izs,ifk,ith,wfld) result(st)
         integer, intent(in) :: iws,izs,ifk,ith
         complex, dimension(:,:,:,:,:), pointer :: wfld
       end function FXop
    end interface !!------------------------------------------
#ifdef DEBUG
    if(wpro%dbg>9) call in(unit)
#endif

    if(wpro%adj) then !! MIGRATION
       call weiift(iws,        ith,wfld)
       st   = FKop(iws,izs,  1,ith,wfld)      !! f-k op
       call weifft(iws,        ith,wfld)
    else              !! MODELLING
       call weiift(iws,        ith,wfld)
       st   = FKop(iws,izs,  1,ith,wfld)      !! f-k op
       call weifft(iws,        ith,wfld)
    end if

#ifdef DEBUG
    if(wpro%dbg>9) call out(unit)
#endif
    st=weiOK
  end function weiop_mwc0
  !----------------------------------------------------------------
  !----------------------------------------------------------------
  integer function weiop_mwc1_init() result(st)
    integer :: er
    character(len=128), parameter :: unit='weiop_mwc1_init'
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif
    !! do nothing
#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    st=weiOK
  end function weiop_mwc1_init
  !----------------------------------------------------------------
  !! 1 reference slowness
  integer function weiop_mwc1(wfld,iws,izs,ith,FKop,FXop) result(st)
    complex, dimension(:,:,:,:,:), pointer :: wfld
    integer, intent(in)       :: iws,izs,ith
    integer                   :: ifk
    character(len=128), parameter :: unit='weiop_mwc1'
    interface !!----------------------------------------------
       integer function FKop(iws,izs,ifk,ith,wfld) result(st)
         integer, intent(in) :: iws,izs,ifk,ith
         complex, dimension(:,:,:,:,:), pointer :: wfld
       end function FKop
       integer function FXop(iws,izs,ifk,ith,wfld) result(st)
         integer, intent(in) :: iws,izs,ifk,ith
         complex, dimension(:,:,:,:,:), pointer :: wfld
       end function FXop
    end interface !!------------------------------------------
#ifdef DEBUG
    if(wpro%dbg>9) call in(unit)
#endif

    if(wpro%adj) then !! MIGRATION
       call weicomp_r(ith)
       st   = FXop(iws,izs,  1,ith,wfld)
       call weiift(iws,        ith,wfld)
       st   = FKop(iws,izs,  1,ith,wfld)
       call weifft(iws,        ith,wfld)

       call weicomp_s(ith)
       st   = FXop(iws,izs,  1,ith,wfld)
       call weiift(iws,        ith,wfld)
       st   = FKop(iws,izs,  1,ith,wfld)
       call weifft(iws,        ith,wfld)
    else              !! MODELLING
       call weicomp_s(ith)
       call weiift(iws,        ith,wfld)
       st   = FKop(iws,izs,  1,ith,wfld)
       call weifft(iws,        ith,wfld)
       st   = FXop(iws,izs,  1,ith,wfld)

       call weicomp_r(ith)
       call weiift(iws,        ith,wfld)
       st   = FKop(iws,izs,  1,ith,wfld)
       call weifft(iws,        ith,wfld)
       st   = FXop(iws,izs,  1,ith,wfld)
    end if

#ifdef DEBUG
    if(wpro%dbg>9) call out(unit)
#endif
    st=weiOK
  end function weiop_mwc1
  !----------------------------------------------------------------
  !! N reference slownesses
  integer function weiop_mwcN_init() result(st)
    integer :: er
    character(len=128), parameter :: unit='weiop_mwcN_init'
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    !! allocate tmp arrays for interpolation of multiple fk ops
    call weimem_add(8,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n*wpar%nomp,unit)
    allocate(tmp1(    wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n,wpar%nomp),stat=er)
    if(er/=0) call seperr('cant allocate tmp1')

    call weimem_add(8,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n*wpar%nomp,unit)
    allocate(tmp2(    wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n,wpar%nomp),stat=er)
    if(er/=0) call seperr('cant allocate tmp2')

    call weimem_add(4,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n*wpar%nomp,unit)
    allocate(imap(    wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n,wpar%nomp),stat=er)
    if(er/=0) call seperr('cant allocate imap')

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    st=weiOK
  end function weiop_mwcN_init
  !----------------------------------------------------------------
  !! N reference slownesses
  integer function weiop_mwcN(wfld,iws,izs,ith,FKop,FXop) result(st)
    complex, dimension(:,:,:,:,:), pointer :: wfld
    integer, intent(in)       :: iws,izs,ith
    integer                   :: ifk
    integer                   :: ic
    character(len=128), parameter :: unit='weiop_mwcN'
    interface !!----------------------------------------------
       integer function FKop(iws,izs,ifk,ith,wfld) result(st)
         integer, intent(in) :: iws,izs,ifk,ith
         complex, dimension(:,:,:,:,:), pointer :: wfld
       end function FKop
       integer function FXop(iws,izs,ifk,ith,wfld) result(st)
         integer, intent(in) :: iws,izs,ifk,ith
         complex, dimension(:,:,:,:,:), pointer :: wfld
       end function FXop
    end interface !!------------------------------------------
#ifdef DEBUG
    if(wpro%dbg>9) call in(unit)
#endif

    if(wpro%adj) then !! MIGRATION
       !----------------------------------------------------------------       
       do ic=1,2
          if(ic==1) then
             call weicomp_s(ith)
          else
             call weicomp_r(ith)
          end if

          tmp2(:,:,:,:,ith) = wfld(:,:,:,:,iws)
          tmp1(:,:,:,:,ith) = (0.,0.)
          do ifk=1,srp(ith)%nr(izs)
             !             where(srp(ith)%si(:,:,:,:,izs) == ifk)
             !                wfld(:,:,:,:,iws) = tmp2(:,:,:,:,ith)
             !             elsewhere
             !                wfld(:,:,:,:,iws) = (0.,0.)
             !             end where

             call wei_imap(izs,ith,ifk)
             wfld(:,:,:,:,iws) = &
             tmp2(:,:,:,:,ith) * imap(:,:,:,:,ith)

             st   = FXop(iws,izs,ifk,ith,wfld)
             call weiift(iws,        ith,wfld)
             st   = FKop(iws,izs,ifk,ith,wfld)
             call weifft(iws,        ith,wfld)

             tmp1(:,:,:,:,ith) = &
             tmp1(:,:,:,:,ith) + &
             wfld(:,:,:,:,iws)

          end do
          wfld(:,:,:,:,iws) = tmp1(:,:,:,:,ith)
       end do
       !----------------------------------------------------------------
    else             !! MODELLING
       !----------------------------------------------------------------
       do ic=1,2
          if(ic==1) then
             call weicomp_r(ith)
          else
             call weicomp_s(ith)
          end if

          tmp1(:,:,:,:,ith) = wfld(:,:,:,:,iws)
          tmp2(:,:,:,:,ith) = 0.
          do ifk=1,srp(ith)%nr(izs)
             wfld(:,:,:,:,iws) = tmp1(:,:,:,:,ith)

             call weiift(iws,        ith,wfld)
             st   = FKop(iws,izs,ifk,ith,wfld)
             call weifft(iws,        ith,wfld)
             st   = FXop(iws,izs,ifk,ith,wfld)

!             where(srp(ith)%si(:,:,:,:,izs) == ifk)
!                tmp2(:,:,:,:,ith) = &
!                tmp2(:,:,:,:,ith) + &
!                wfld(:,:,:,:,iws)
!             end where

             call wei_imap(izs,ith,ifk)
             tmp2(:,:,:,:,ith) = &
             tmp2(:,:,:,:,ith) + &
             wfld(:,:,:,:,iws) * imap(:,:,:,:,ith)

          end do
          wfld(:,:,:,:,iws) = tmp2(:,:,:,:,ith)
       end do
    end if

#ifdef DEBUG
    if(wpro%dbg>9) call out(unit)
#endif
    st=weiOK
  end function weiop_mwcN
  !----------------------------------------------------------------  
  subroutine wei_imap(izs,ith,ifk)
    integer, intent(in) :: izs,ith,ifk
    imap(:,:,:,:,ith)=0.
    where( srp(ith)%sw(:,:,:,:,izs) > ifk-1.+ epsilon(1.) .and. srp(ith)%sw(:,:,:,:,izs) < ifk+1. )
       imap(:,:,:,:,ith) = 1.-abs(srp(ith)%sw(:,:,:,:,izs)-ifk)
    end where
  end subroutine wei_imap
  !----------------------------------------------------------------
end module weiop_mwc


!! multiple fk ops
!!
!!       D(z)
!!       |
!!       v
!!       T1
!!
!!    D1 D2 D3 ...
!!
!!    S1 S2 S3 ...
!!     \ | /
!!       T2
!!       |
!!       v
!!       D(z+dz)



!! Multi-reference slowness operator scheme

!! K = reference slowness mask
!! P = phase shift operator  

!! modeling
!! 
!! Wz = [ K1 K2 K3] [ P1 ] Wo = K P  Wo
!!                  [ P2 ]
!!                  [ P3 ]

!! migration
!!
!! Wo = [ P1 P2 P3] [ P1 ] Wo = P'K' Wz
!!                  [ P2 ]
!!                  [ P3 ]
