module weiop_ryt
  !! Rytov WEMVA

  use wei_util
  use wei_pro
  use wei_par
  use wei_S
  use wei_kmap
  use wei_bosc
  !---------------
  implicit none

  complex, dimension(:,:,:,:,:), pointer, private :: dwsou,dwrec
  complex, dimension(:,:,:,:,:), pointer, private :: dssou,dsrec

  real,    dimension(:,:,:,:),   pointer, private :: kss,krr
contains
  !----------------------------------------------------------------
  integer function weiop_ryt1_init() result(st)
    integer :: imx,ihx, imy,ihy
    integer :: er
    character(len=128), parameter :: unit='weiop_ryt1_init'
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

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

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

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

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

    !! ks(mx,my,hx,hy)
    call weimem_add(4,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n,unit)
    allocate(kss(     wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n),stat=er)
    if(er/=0) call seperr('cant allocate kss')

    !! kr(mx,my,hx,hy)
    call weimem_add(4,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n,unit)
    allocate(krr(     wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n),stat=er)
    if(er/=0) call seperr('cant allocate krr')

    do ihy=1,wpro%ahy%n
       do ihx=1,wpro%ahx%n
          do imy=1,wpro%amy%n
             do imx=1,wpro%amx%n
                kss(imx,imy,ihx,ihy) = &
                sqrt(kmap%ksx(imx,ihx)**2   &
                +    kmap%ksy(imy,ihy)**2)

                krr(imx,imy,ihx,ihy) = &
                sqrt(kmap%krx(imx,ihx)**2   &
                +    kmap%kry(imy,ihy)**2)
             end do
          end do
       end do
    end do

    !! init scattering operator
    st=weibosc_init()

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    st=weiOK
  end function weiop_ryt1_init
  !----------------------------------------------------------------
  integer function weiop_ryt1(bwfld,dslow,dwfld,iws,izs,ith) result(st)
    complex, dimension(:,:,:,:,:), pointer :: bwfld,dslow,dwfld
    integer, intent(in)       :: iws,izs,ith
    real                      :: w,kos,kor
    real                      :: sos,sor
    integer :: i1,i2,i3,i4
    real    :: eps
    character(len=128), parameter :: unit='weiop_ryt1'
#ifdef DEBUG
    if(wpro%dbg>9) call in(unit)
#endif
    w  = wpro%bw%o + (iws-1)*wpro%bw%d

    sos = sum(slo%rs_s(izs,:))/slo%nr_s(izs)
    sor = sum(slo%rs_r(izs,:))/slo%nr_r(izs)

    kos = 2*w * sos
    kor = 2*w * sos

    eps = 0.1*maxval(cabs(bwfld(:,:,:,:,iws)))

    if(wpro%adj) then !! velocity analysis
#ifdef DEBUG
       if(wpro%dbg>9) call  in('dS<-dW')
#endif
       dwsou(:,:,:,:,ith) = dwfld(:,:,:,:,iws)
       dwrec(:,:,:,:,ith) = dwfld(:,:,:,:,iws)

       where(cabs(bwfld(:,:,:,:,iws))>eps)
          dwsou(:,:,:,:,ith) = &
          dwsou(:,:,:,:,ith) /      (bwfld(:,:,:,:,iws))
          dwrec(:,:,:,:,ith) = &
          dwrec(:,:,:,:,ith) /      (bwfld(:,:,:,:,iws))
!       elsewhere 
!          dwsou(:,:,:,:,ith)=0.
!          dwrec(:,:,:,:,ith)=0.
       end where

       st= weibosc_w2s(bwfld,iws,ith,dssou,dwsou,kss,kos)
       st= weibosc_w2s(bwfld,iws,ith,dsrec,dwrec,krr,kor)

       !! dS += (dSs + dSr)
       !$OMP ATOMIC
       dslow(:,:,:,:,izs) = &
       dslow(:,:,:,:,izs) + &
       dssou(:,:,:,:,ith) + &
       dsrec(:,:,:,:,ith) 

#ifdef DEBUG
       if(wpro%dbg>9) call out('dS<-dW')
#endif
    else !! scaterring
#ifdef DEBUG
       if(wpro%dbg>9) call  in('dS->dW')
#endif

       dssou(:,:,:,:,ith) = dslow(:,:,:,:,izs)
       dsrec(:,:,:,:,ith) = dslow(:,:,:,:,izs)

       st= weibosc_s2w(bwfld,iws,ith,dssou,dwsou,kss,kos)
       st= weibosc_s2w(bwfld,iws,ith,dsrec,dwrec,krr,kor)

       where(cabs(bwfld(:,:,:,:,iws))>eps)
          dwsou(:,:,:,:,ith) = &
          dwsou(:,:,:,:,ith) / conjg(bwfld(:,:,:,:,iws))
          dwrec(:,:,:,:,ith) = &
          dwrec(:,:,:,:,ith) / conjg(bwfld(:,:,:,:,iws))
!       elsewhere
!          dwsou(:,:,:,:,ith)=0.
!          dwrec(:,:,:,:,ith)=0.
       end where

!       !! dW += (dWs + dWr)
!       dwfld(:,:,:,:,iws) = &
!       dwfld(:,:,:,:,iws) + &
!       dwsou(:,:,:,:,ith) + &
!       dwrec(:,:,:,:,ith)

       dwfld(:,:,:,:,iws) = &
       cexp(clog(dwfld(:,:,:,:,iws)) + &
       dwsou(:,:,:,:,ith) + &
       dwrec(:,:,:,:,ith))

#ifdef DEBUG
       if(wpro%dbg>9) call out('dS->dW')
#endif
    end if

#ifdef DEBUG
    if(wpro%dbg>9) call out(unit)
#endif
    st=weiOK
  end function weiop_ryt1
  !----------------------------------------------------------------
end module weiop_ryt
