module wei_mva
  !! WEI migration velocity analysis
 
  use wei_util
  use wei_pro
  use wei_par
  use wei_tag
  use wei_comp

  use wei_S
  use wei_R
  use wei_D
  use wei_W
  use wei_dS

  use wei_fft
  use wei_kmap
  use wei_xmap

  use wei_image
  use wei_kxmig
  use wei_kxmva

  use wei_attr
  !---------------
  implicit none

  logical, private :: wIO

  type(weslow), private :: SS
  type(wedata), private :: DD 
  type(weimag), private :: RR
  
  type(wedata), public  :: bW
  type(wewfld), private :: WW

  type(wedata), private :: dW

  type(weslow), private :: dS
  type(weimag), private :: dR
contains
  !----------------------------------------------------------------
  subroutine weimva_init(SLin,WCin,SCin,FKin,FXin,IGin,in7,in8,in9)
    integer,external,optional ::                       in7,in8,in9
    integer :: st,er
    character(len=128) :: name
#ifdef DEBUG
    character(len=128), parameter :: unit='weimva_init'
#endif
    interface !!-------------------------
       integer function SLin() result(st)
       end function SLin
       integer function WCin() result(st)
       end function WCin
       integer function SCin() result(st)
       end function SCin
       integer function FKin() result(st)
       end function FKin
       integer function FXin() result(st)
       end function FXin
       integer function IGin() result(st)
       end function IGin
    end interface !!---------------------

#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    wpro%mod_ = wtag%dS
    wpro%dat_ = wtag%dR

    bW%tag="bW"; call weiD_aux_init(bW)
    dW%tag="dW"; call weiD_aux_init(dW)

    !! initialize objects
    SS%tag=wtag%S;  call weiS_init( SS)
    dS%tag=wtag%dS; call weidS_init(dS)
    dR%tag=wtag%dR; call weiR_init( dR)

    !! flag whether I/O on wavefield
    wIO=.false.
    if(wpro%operator(1:5)=='dSwdR') wIO=.true.
    if(wIO) then
       WW%tag=wtag%W; call weiW_init(WW,bW%bin)
    else
       RR%tag=wtag%R; call weiR_init(RR)
       DD%tag=wtag%D; call weiD_init(DD)
    end if

    !! convert frequency to radians
    call weipro_radians()

    call weifft_init()    !! init FFT
    call weikmap_init()   !! init K-map
    call weixmap_init()   !! init X-map
    call weiimage_init()  !! init imaging

    !! allocate slowness arrays for FK/FX ops
    call  weiS_fkfx_allocate()
    call weidS_fkfx_allocate()

    !! init source/receiver pointers
    call weicomp_init()

    !! check if offset is consistent with ph
    if(wpro%pCIG) call weipro_check_hp()

    !! init operators
    st=SLin()  !! SL op init
    st=WCin()  !! WC op init
    st=SCin()  !! SC op init
    st=FKin()  !! FK op init
    st=FXin()  !! FX op init
    st=IGin()  !! IG op init

    !! spread slowness on the nodes
    call weiS_mpicopy(SS)

    !! spread W or R+D to the nodes
    if(wIO) then
       if(wpro%wTRA) call weiW_mpisprd(WW,bW%bin)
    else
       call weiR_mpicopy(RR)
       call weiD_mpisprd(DD)
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weimva_init
  !----------------------------------------------------------------
  !----------------------------------------------------------------
  subroutine weimva_outnull()
#ifdef DEBUG
    character(len=128), parameter :: unit='weimva_outnull'
    if(wpro%dbg>0) call in(unit)
#endif

    if(wpro%adj) then !! dR->dS
       call weiS_assign(dS,cmplx(0.))
    else              !! dR<-dS
       call weiR_assign(dR,cmplx(0.))
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weimva_outnull
  !----------------------------------------------------------------
  !----------------------------------------------------------------
  integer function weimva(adj,add,mod_,dat_,SLop,WCop,SCop,FKop,FXop,IGop,op7,op8,op9)
    integer,external,optional      ::       SLop,WCop,SCop,FKop,FXop,IGop,op7,op8,op9
    logical,            intent(in) :: adj,add
    character(len=128), intent(in) :: mod_,dat_
    integer :: izb,iwb
    integer :: st
#ifdef DEBUG
    character(len=128), parameter :: unit='weimva'
    if(wpro%dbg>0) call in(unit)
#endif

    !! reset tag variables
    call weipro_reset(adj,add,mod_,dat_)
    if(wpar%impi==0 .and. wpro%verb) call weipro_follow() !! where am I?

    !! accomodate tag name changes
    !! this is only needed on the master
    !! since the files on the nodes are temporary
    !! and going to be filled next
    if(wpar%impi==0) then
       dS%tag = trim(wpro%mod_)
       dR%tag = trim(wpro%dat_)
    end if

    !! zero the output if add is off
    if(.not. wpro%add) call weimva_outnull()

    !! setup the adjoint sign
    if(wpro%adj) then
       wpro%opsign = weiB2T !! dS <- dR (from bottom to top)
    else
       wpro%opsign = weiT2B !! dS -> dR (from top to bottom)
    end if

    if(wpar%nmpi>1) then
       !! HERE DISTRIBUTE DATA FROM MASTER
       !! distribute dS & dR
       if(wpro%adj) then
          call weiS_mpizero(dS)         !! dSi = 0
          call weiR_mpicopy(dR)         !! dRi = dR
       else
          call weiS_mpicopy(dS)         !! dSi = dS
          call weiR_mpizero(dR)         !! dRi = 0
       end if
    end if

    write(0,'(a,i3,a)') weitab,wpar%impi," begin"
    !===========================================
    if(wpro%outw) then
       weimva=weimvaof(SLop,WCop,SCop,FKop,FXop,IGop)
    else 
       weimva=weimvaoz(SLop,WCop,SCop,FKop,FXop,IGop)
    end if
    !===========================================
    write(0,'(a,i3,a)') weitab,wpar%impi," end "

    if(wpar%nmpi>1) then
       !! HERE COLLECT DATA FROM NODES
       !! collect dS & dR
       if(wpro%adj) then
          call weiS_mpisumm(dS)         !! dS = dS + dSi
       else
          call weiR_mpisumm(dR)         !! dR = dR + dRi
       end if
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    weimva=weiOK
  end function weimva
  !----------------------------------------------------------------
  integer function weimvaof(SLop,WCop,SCop,FKop,FXop,IGop,op7,op8,op9)
    integer,external,optional :: WCop,SCop,FKop,FXop,IGop,op7,op8,op9
    integer :: izb,iwb
    integer :: st
#ifdef DEBUG
    character(len=128), parameter :: unit='weimvaof'
#endif
    interface !!----------------------------------------------
       integer function SLop(X) result(st)
         use wei_type
         type(weslow), intent(inout) :: X
       end function SLop
    end interface !!------------------------------------------
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    if(wpro%adj) then !! ADJOINT operator
       !----------------------------------------------------------------
       !! march from bottom to top (pairs w/ modelling)
#ifdef DEBUG
       if(wpro%dbg>0) call  in('dS<-dR')
#endif

       call weiD_assign(bW,cmplx(0.))  !! start with no data
       call weiD_assign(dW,cmplx(0.))

       wbA: do iwb=1,wpro%nwb
          wpro%bw%o = wpro%nw%o + (iwb-1)*(wpro%bw%n*wpro%bw%d)

          if(wIO) then
             call weiD_read(dW,iwb)
          else
             call weiD_read(bW,iwb)
             call weiD_read(dW,iwb)
          end if

          zbA: do izb=wpro%nzb,1,-1
             if(wpro%verb) call weipro_followup(izb,iwb)

             if(.not. wIO) then
                call weiR_read(RR,izb)  !! read image block
             end if
             call weiS_read(SS,izb)     !! read a slowness block
             st = SLop(SS)              !! setup sS, sR, s(z)

             call weiR_read(dR,izb)
             call weiS_read(dS,izb)

             dsl%ds_s=0.
             dsl%ds_r=0.
             if(wIO) then
                call wmvaWW(izb,iwb,WCop,SCop,FKop,FXop,IGop) !! read W
             else
                call wmvaRD(izb,iwb,WCop,SCop,FKop,FXop,IGop) !! trash W
             end if

             call weidS_pba2bin(dS)
             call weiS_write(dS,izb)
          end do zbA
       end do wbA

#ifdef DEBUG
       if(wpro%dbg>0) call out('dS<-dR')
#endif
       !----------------------------------------------------------------
    else             !! FORWARD operator
       !----------------------------------------------------------------
       !! march from top to bottom (pairs w/ migration)
#ifdef DEBUG
       if(wpro%dbg>0) call  in('dS->dR')
#endif

       if(.not. wIO) then
          call weiD_copy(DD,bW)      !! save the data at the top
       end if
       call weiD_assign(dW,cmplx(0.))

       wbF: do iwb=1,wpro%nwb
          wpro%bw%o = wpro%nw%o + (iwb-1)*(wpro%bw%n*wpro%bw%d)

          if(wIO) then
             call weiD_read(dW,iwb)
          else
             call weiD_read(bW,iwb)
             call weiD_read(dW,iwb)
          end if

          zbF: do izb=1,wpro%nzb
             if(wpro%verb) call weipro_followup(izb,iwb)

             call weiS_read(SS,izb)      !! read a slowness block
             st = SLop(SS)               !! setup sS,sR, s(z)

             call weiR_read(dR,izb)
             call weiS_read(dS,izb)
             call weidS_bin2pba(dS)

             if(wIO) then
                call wmvaWW(izb,iwb,WCop,SCop,FKop,FXop,IGop) !! read W
             else
                call wmvaRD(izb,iwb,WCop,SCop,FKop,FXop,IGop) !! trash W
             end if

             call weiR_write(dR,izb)
          end do zbF
       end do wbF

#ifdef DEBUG
       if(wpro%dbg>0) call out('dS->dR')
#endif
       !----------------------------------------------------------------
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    weimvaof=weiOK
  end function weimvaof
  !----------------------------------------------------------------
  integer function weimvaoz(SLop,WCop,SCop,FKop,FXop,IGop,op7,op8,op9)
    integer,external,optional :: WCop,SCop,FKop,FXop,IGop,op7,op8,op9
    integer :: izb,iwb
    integer :: st
#ifdef DEBUG
    character(len=128), parameter :: unit='weimvaoz'
#endif
    interface !!----------------------------------------------
       integer function SLop(X) result(st)
         use wei_type
         type(weslow), intent(inout) :: X
       end function SLop
    end interface !!------------------------------------------
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    if(wpro%adj) then !! ADJOINT operator
       !----------------------------------------------------------------
       !! march from bottom to top (pairs w/ modelling)
#ifdef DEBUG
       if(wpro%dbg>0) call  in('dS<-dR')
#endif

       call weiD_assign(bW,cmplx(0.))  !! start with no data
       call weiD_assign(dW,cmplx(0.))

       zbA: do izb=wpro%nzb,1,-1
          if(.not. wIO) then
             call weiR_read(RR,izb)    !! read image block
          end if
          call weiS_read(SS,izb)       !! read a slowness block
          st = SLop(SS)                !! setup sS, sR, s(z)

          call weiR_read(dR,izb)
          call weiS_read(dS,izb)

          dsl%ds_s=0.
          dsl%ds_r=0.
          wbA: do iwb=1,wpro%nwb
             if(wpro%verb) call weipro_followup(izb,iwb)
             wpro%bw%o = wpro%nw%o + (iwb-1)*(wpro%bw%n*wpro%bw%d)

             if(wIO) then
                call weiD_read(dW,iwb)  !! read  @ bottom
                call wmvaWW(izb,iwb,WCop,SCop,FKop,FXop,IGop) !! read W
                call weiD_write(dW,iwb) !! write @ top
             else
                call weiD_read(bW,iwb)
                call weiD_read(dW,iwb)  !! read  @ bottom
                call wmvaRD(izb,iwb,WCop,SCop,FKop,FXop,IGop) !! trash W
                call weiD_write(dW,iwb) !! write @ top
                call weiD_write(bW,iwb)
             end if

          end do wbA

          call weidS_pba2bin(dS)
          call weiS_write(dS,izb)
       end do zbA

#ifdef DEBUG
       if(wpro%dbg>0) call out('dS<-dR')
#endif
       !----------------------------------------------------------------
    else             !! FORWARD operator
       !----------------------------------------------------------------
       !! march from top to bottom (pairs w/ migration)
#ifdef DEBUG
       if(wpro%dbg>0) call  in('dS->dR')
#endif

       if(.not. wIO) then
          call weiD_copy(DD,bW)      !! save the data at the top
       end if
       call weiD_assign(dW,cmplx(0.))

       zbF: do izb=1,wpro%nzb
          call weiS_read(SS,izb)     !! read a slowness block
          st = SLop(SS)              !! setup sS,sR, s(z)

          call weiR_read(dR,izb)
          call weiS_read(dS,izb)
          call weidS_bin2pba(dS)

          wbF: do iwb=1,wpro%nwb
             if(wpro%verb) call weipro_followup(izb,iwb)
             wpro%bw%o = wpro%nw%o + (iwb-1)*(wpro%bw%n*wpro%bw%d)

             if(wIO) then
                call weiD_read(dW,iwb) !! read  @ top
                call wmvaWW(izb,iwb,WCop,SCop,FKop,FXop,IGop) !! read W
                call weiD_write(dW,iwb)!! write @ bottom
             else
                call weiD_read(bW,iwb)
                call weiD_read(dW,iwb) !! read  @ top
                call wmvaRD(izb,iwb,WCop,SCop,FKop,FXop,IGop) !! trash W
                call weiD_write(dW,iwb)!! write @ bottom
                call weiD_write(bW,iwb)
             end if
          end do wbF

          call weiR_write(dR,izb)
       end do zbF

#ifdef DEBUG
       if(wpro%dbg>0) call out('dS->dR')
#endif
       !----------------------------------------------------------------
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    weimvaoz=weiOK
  end function weimvaoz
  !----------------------------------------------------------------
  subroutine wmvaRD(izb,iwb,WCop,SCop,FKop,FXop,IGop)
    integer, intent(in)       :: izb,iwb
    integer,external,optional :: WCop,SCop,FKop,FXop,IGop
    integer :: izs,iws,ith
#ifdef DEBUG
    character(len=128), parameter :: unit='wmvaRD'
    if(wpro%dbg>0) call in(unit)
#endif

    if(wpro%adj) then !! ADJOINT dS<-dR
       !$OMP  PARALLEL
       !$OMP  DO &
       !$OMP& SCHEDULE(DYNAMIC) &
       !$OMP& PRIVATE(iws,izs,ith)
       wsA: do iws=1,wpro%nws
          ith = omp_get_thread_num()+1
          zsA: do izs=wpro%nzs,1,-1

             !! UC W to current depth level
             call weikxmig(bW%bin,iws,izs,ith,WCop,FKop,FXop)
             !! inject R at current level
             call weiimgR2W(izs,iws,ith,RR%bin,bW%bin, IGop)

             !! UC dW to current depth level
             call weikxmig(dW%bin,iws,izs,ith,WCop,FKop,FXop)
             !! inject dR at current level
             call weiimgR2W(izs,iws,ith, dR%bin,dW%bin, IGop)
             !! extract dS from dW
             call weikxmva(bW%bin,dW%bin,dsl%ds_s,dsl%ds_r,iws,izs,ith,SCop)

          end do zsA
       end do wsA
       !$OMP END DO
       !$OMP END PARALLEL
    else             !! FORWARD dS->dR
       !$OMP  PARALLEL
       !$OMP  DO &
       !$OMP& SCHEDULE(DYNAMIC) &
       !$OMP& PRIVATE(iws,izs,ith)
       wsF: do iws=1,wpro%nws
          ith = omp_get_thread_num()+1
          zsF: do izs=1,wpro%nzs

             !! scatter dS in dW
             call weikxmva(bW%bin,dW%bin,dsl%ds_s,dsl%ds_r,iws,izs,ith,SCop)
             !! sum dW to dR
             call weiimgW2R(izs,iws,ith, dR%bin,dW%bin,IGop)
             !! DC dW to next deeper level
             call weikxmig(dW%bin,iws,izs,ith,WCop,FKop,FXop)

             !! DC  W to next deeper level
             call weikxmig(bW%bin,iws,izs,ith,WCop,FKop,FXop)

          end do zsF
       end do wsF
       !$OMP END DO
       !$OMP END PARALLEL
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine wmvaRD
  !----------------------------------------------------------------
  !----------------------------------------------------------------
  subroutine wmvaWW(izb,iwb,WCop,SCop,FKop,FXop,IGop)
    integer, intent(in)       :: izb,iwb
    integer,external,optional :: WCop,SCop,FKop,FXop,IGop
    integer :: izs,iws,ith
#ifdef DEBUG
    character(len=128), parameter :: unit='wmvaWW'
    if(wpro%dbg>0) call in(unit)
#endif

    if(wpro%adj) then !! ADJOINT dS<-dR
       !$OMP  PARALLEL
       !$OMP  DO &
       !$OMP& SCHEDULE(DYNAMIC) &
       !$OMP& PRIVATE(iws,izs,ith)
       wsA: do iws=1,wpro%nws
          ith = omp_get_thread_num()+1
          zsA: do izs=wpro%nzs,1,-1

             !$OMP CRITICAL
             call weiW_read(WW,bW%bin,izb,izs,iwb,iws)
             !$OMP END CRITICAL

             !! UC dW to current depth level
             call weikxmig(dW%bin,iws,izs,ith,WCop,FKop,FXop)
             !! inject dR at current level
             call weiimgR2W(izs,iws,ith, dR%bin,dW%bin, IGop)
             !! extract dS from dW
             call weikxmva(bW%bin,dW%bin,dsl%ds_s,dsl%ds_r,iws,izs,ith,SCop)

          end do zsA
       end do wsA
       !$OMP END DO
       !$OMP END PARALLEL
    else             !! FORWARD dS->dR
       !$OMP  PARALLEL
       !$OMP  DO &
       !$OMP& SCHEDULE(DYNAMIC) &
       !$OMP& PRIVATE(iws,izs,ith)
       wsF: do iws=1,wpro%nws
          ith = omp_get_thread_num()+1
          zsF: do izs=1,wpro%nzs

             !$OMP CRITICAL
             call weiW_read(WW,bW%bin,izb,izs,iwb,iws)
             !$OMP END CRITICAL

             !! scatter dS in dW
             call weikxmva(bW%bin,dW%bin,dsl%ds_s,dsl%ds_r,iws,izs,ith,SCop)
             !! sum dW to dR
             call weiimgW2R(izs,iws,ith, dR%bin,dW%bin,IGop)
             !! DC dW to next deeper level
             call weikxmig(dW%bin,iws,izs,ith,WCop,FKop,FXop)

          end do zsF
       end do wsF
       !$OMP END DO
       !$OMP END PARALLEL
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine wmvaWW
  !----------------------------------------------------------------
  !----------------------------------------------------------------
end module wei_mva



!!     FOR          ADJ
!!   dS    D     dS
!!   |     |     ^     ^
!!   |     |     |     |
!!   | <-W-|     | <-W-|
!!   |     |     |     |
!!   v     v     |     R
!!   dR          dR


!! FORWARD
!!
!!  z  :   U -> dS ==> dU -> dR
!!         |           |
!!         v           v
!!  z+1:   U           dU 
!!
!!


!! ADJOINT
!!
!!
!!    R -> U ==> dS <== dU <- dR
!!         ^            ^
!!         |            |
!! z+1:    U            dU
!!


!! How to get the phase perturbation?

!! version 1 - incorrect
!! Rb 
!!    -> dR -> dW -> dF -unwrap -> dFu 
!! Rp
!!

!! version 2 - correct
!! Rb -> Wb -> Fb -unwrap-> Fub
!!                              -> dFu
!! Rp -> Wp -> Fp -unwrap-> Fup 
!!





