module wei_mig
  !! WEI migration

  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_fft
  use wei_kmap
  use wei_xmap

  use wei_image
  use wei_kxmig
  use wei_dump
  !---------------
  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
contains
  !----------------------------------------------------------------
  subroutine weimig_init(SLin,WCin,FKin,FXin,IGin,in6,in7,in8,in9)
    integer,external,optional ::                  in6,in7,in8,in9
    integer            :: st,er
    character(len=128) :: name
#ifdef DEBUG
    character(len=128), parameter :: unit='weimig_init'
#endif
    interface !!-------------------------
       integer function  SLin() result(st)
       end function SLin
       integer function  WCin() result(st)
       end function WCin
       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%R
    wpro%dat_ = wtag%D

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

    !! initialize objects
    SS%tag=wtag%S; call weiS_init(SS)
    DD%tag=wtag%D; call weiD_init(DD)
    RR%tag=wtag%R; call weiR_init(RR)

    wIO=.false.
    if(wpro%operator(1:5)=='.Rw.D') wIO=.true.
    if(wIO) then
       WW%tag=wtag%W; call weiW_init(WW,bW%bin)
    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()

    !! 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=FKin()     !! FK op init
    st=FXin()     !! FX op init
    st=IGin()     !! FX op init

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

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

    if(wpro%adj) then !! R<-D
       call weiR_assign(RR,(0.,0.))
    else              !! R->D
       call weiD_assign(DD,(0.,0.))
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weimig_outnull
  !----------------------------------------------------------------
  !----------------------------------------------------------------
  integer function weimig(adj,add,mod_,dat_,SLop,WCop,FKop,FXop,IGop,op6,op7,op8,op9)
    integer,external,optional      ::       SLop,WCop,FKop,FXop,IGop,op6,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='weimig'
    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
       RR%tag = trim(wpro%mod_)
       DD%tag = trim(wpro%dat_)
    end if

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

    !! setup the adjoint sign
    if(wpro%adj) then
       wpro%opsign = weiT2B
    else
       wpro%opsign = weiB2T
    end if

    if(wpar%nmpi>1) then
       !! HERE DISTRIBUTE DATA FROM MASTER
       !! distribute RR & DD
       if(wpro%adj) then
          call weiR_mpizero(RR)              !! Ri = 0
          call weiD_mpisprd(DD)              !! D -> [Do D1 ... Dn]
       else
          call weiR_mpicopy(RR)              !! Ri = R
          call weiD_mpizero(DD)              !! Di = 0
       end if
       if(wIO) then
          if(wpro%wTRA) call weiW_mpizero(WW,bW%bin)  !! Wi = 0
       end if
    end if

    write(0,'(a,i3,a)') weitab,wpar%impi," begin"
    !===========================================
    if(wpro%outw) then
       weimig=weimigof(SLop,WCop,FKop,FXop,IGop)
    else 
       weimig=weimigoz(SLop,WCop,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 RR & DD & bW
       if(wpro%adj) then
          call weiR_mpisumm(RR)              !! R = R + Ri
       else
          call weiD_mpigthr(DD)              !! D <- [Do D1 ... Dn]
       end if
       if(wIO) then
          if(wpro%wTRA) call weiW_mpigthr(WW,bW%bin)    !! W <- [Wo W1 ... Wn]
       end if
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    weimig=weiOK
  end function weimig
  !----------------------------------------------------------
  integer function weimigof(SLop,WCop,FKop,FXop,IGop,op6,op7,op8,op9)
    integer,external,optional :: WCop,FKop,FXop,IGop,op6,op7,op8,op9
    integer :: izb,iwb
    integer :: st
#ifdef DEBUG
    character(len=128), parameter :: unit='weimigof'
#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 !! MIGRATION
#ifdef DEBUG
       if(wpro%dbg>0) call in('migration')
#endif
       !----------------------------------------------------------
       call weiD_copy(DD,bW)             !! copy the data at the top
       wbA: do iwb=1,wpro%nwb
          wpro%bw%o = wpro%nw%o + (iwb-1)*(wpro%bw%n*wpro%bw%d)
          call weiD_read(bW,iwb)         !! read  data at block top
          zbA: do izb=1,wpro%nzb       
             if(wpro%verb) call weipro_followup(izb,iwb)
             call weiR_read(RR,izb)      !! read image block
             call weiS_read(SS,izb)      !! read a slowness block
             st = SLop(SS)               !! setup sS, sR, s(z)
             if(wIO) then                !! DC one block
                call wmigWW(izb,iwb,WCop,FKop,FXop,IGop) !! write W
             else
                call wmigRD(izb,iwb,WCop,FKop,FXop,IGop) !! trash W
             end if
             call weiR_write(RR,izb)     !! write image block
          end do zbA
          call weiD_write(bW,iwb)        !! write data at block bottom
       end do wbA
       !----------------------------------------------------------
#ifdef DEBUG
       if(wpro%dbg>0) call out('migration')
#endif
    else         !! MODELING
#ifdef DEBUG
       if(wpro%dbg>0) call in('modeling')
#endif
       !----------------------------------------------------------
       call weiD_assign(bW,cmplx(0.))    !! start with no data
       wbF: do iwb=1,wpro%nwb
          wpro%bw%o = wpro%nw%o + (iwb-1)*(wpro%bw%n*wpro%bw%d)
          call weiD_read(bW,iwb)         !! read  data at block bottom
          zbF: do izb=wpro%nzb,1,-1    
             if(wpro%verb) call weipro_followup(izb,iwb)
             call weiR_read(RR,izb)      !! read image block
             call weiS_read(SS,izb)      !! read a slowness block
             st = SLop(SS)               !! setup sS,sR, s(z)
             if(wIO) then                !! DC one block
                call wmigWW(izb,iwb,WCop,FKop,FXop,IGop) !! write W
             else
                call wmigRD(izb,iwb,WCop,FKop,FXop,IGop) !! trash W
             end if
          end do zbF
          call weiD_write(bW,iwb)        !! write data at block top
       end do wbF
       call weiD_add(bW,DD)              !! add U.C. wfld to data
       !----------------------------------------------------------
#ifdef DEBUG
       if(wpro%dbg>0) call out('modeling')
#endif
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    weimigof=weiOK
  end function weimigof
  !----------------------------------------------------------
  integer function weimigoz(SLop,WCop,FKop,FXop,IGop,op6,op7,op8,op9)
    integer,external,optional :: WCop,FKop,FXop,IGop,op6,op7,op8,op9
    integer :: izb,iwb
    integer :: st
#ifdef DEBUG
    character(len=128), parameter :: unit='weimigoz'
#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 !! MIGRATION
#ifdef DEBUG
       if(wpro%dbg>0) call in('migration')
#endif
       !----------------------------------------------------------
       call weiD_copy(DD,bW)             !! copy the data at the top
       zbA: do izb=1,wpro%nzb        
          call weiR_read(RR,izb)         !! read image block
          call weiS_read(SS,izb)         !! read a slowness block
          st = SLop(SS)                  !! setup sS, sR, s(z)
          wbA: do iwb=1,wpro%nwb
             wpro%bw%o = wpro%nw%o + (iwb-1)*(wpro%bw%n*wpro%bw%d)
             if(wpro%verb) call weipro_followup(izb,iwb)
             call weiD_read(bW,iwb)      !! read  data at block top
             if(wIO) then                !! DC one block
                call wmigWW(izb,iwb,WCop,FKop,FXop,IGop) !! write W
             else
                call wmigRD(izb,iwb,WCop,FKop,FXop,IGop) !! trash W
             end if
             call weiD_write(bW,iwb)     !! write data at block bottom
          end do wbA
          call weiR_write(RR,izb)        !! write image block
       end do zbA
       !----------------------------------------------------------
#ifdef DEBUG
       if(wpro%dbg>0) call out('migration')
#endif
    else         !! MODELING
#ifdef DEBUG
       if(wpro%dbg>0) call in('modeling')
#endif
       !----------------------------------------------------------
       call weiD_assign(bW,cmplx(0.))    !! start with no data
       zbF: do izb=wpro%nzb,1,-1
          call weiR_read(RR,izb)         !! read image block
          call weiS_read(SS,izb)         !! read a slowness block
          st = SLop(SS)                  !! setup sS,sR, s(z)
          wbF: do iwb=1,wpro%nwb
             wpro%bw%o = wpro%nw%o + (iwb-1)*(wpro%bw%n*wpro%bw%d)
             if(wpro%verb) call weipro_followup(izb,iwb)
             call weiD_read(bW,iwb)      !! read  data at block bottom
             if(wIO) then                !! DC one block
                call wmigWW(izb,iwb,WCop,FKop,FXop,IGop) !! write W
             else
                call wmigRD(izb,iwb,WCop,FKop,FXop,IGop) !! trash W
             end if
             call weiD_write(bW,iwb)     !! write data at block top
          end do wbF
       end do zbF
       call weiD_add(bW,DD)              !! add U.C. wfld to data
       !----------------------------------------------------------
#ifdef DEBUG
       if(wpro%dbg>0) call out('modeling')
#endif
    end if
#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
    weimigoz=weiOK
  end function weimigoz
  !---------------------------------------------------------------
  !! don't write the wavefield
  subroutine wmigRD(izb,iwb,WCop,FKop,FXop,IGop)
    integer, intent(in)       :: izb,iwb
    integer,external,optional :: WCop,FKop,FXop,IGop
    integer :: izs,iws,ith
#ifdef DEBUG
    character(len=128), parameter :: unit='wmigRD'
    if(wpro%dbg>0) call in(unit)
#endif

    if(wpro%adj) then !! MIGRATION
       !$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=1,wpro%nzs

             call weiimgW2R(izs,iws,ith,RR%bin,bW%bin,IGop)
             call weikxmig(bW%bin,iws,izs,ith,WCop,FKop,FXop)
             
          end do zsA
       end do wsA
       !$OMP END DO
       !$OMP END PARALLEL
    else         !! MODELLING
       !$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=wpro%nzs,1,-1
             call weikxmig(bW%bin,iws,izs,ith,WCop,FKop,FXop)
             call weiimgR2W(izs,iws,ith, RR%bin,bW%bin,IGop)
          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 wmigRD
  !----------------------------------------------------------------
  !! write the wavefield
  subroutine wmigWW(izb,iwb,WCop,FKop,FXop,IGop)
    integer, intent(in)       :: izb,iwb
    integer,external,optional :: WCop,FKop,FXop,IGop
    integer :: izs,iws,ith
#ifdef DEBUG
    character(len=128), parameter :: unit='wmigWW'
    if(wpro%dbg>0) call in(unit)
#endif

    if(wpro%adj) then !! MIGRATION
       !$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=1,wpro%nzs

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

             call weiimgW2R(izs,iws,ith, RR%bin,bW%bin,IGop)
             call weikxmig(bW%bin,iws,izs,ith,WCop,FKop,FXop)

          end do zsA
       end do wsA
       !$OMP END DO
       !$OMP END PARALLEL
    else         !! MODELLING
       !$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=wpro%nzs,1,-1

             call weikxmig(bW%bin,iws,izs,ith,WCop,FKop,FXop)
             call weiimgR2W(izs,iws,ith, RR%bin,bW%bin,IGop)

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

          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 wmigWW
  !----------------------------------------------------------
end module wei_mig


!!   ADJ         FOR
!!
!! D1 > R1     D1 > R1
!! v           ^
!! D2 > R2     D2 > R2
!! v           ^
!! D3 > R3     D3 > R3
!! v           ^
!! D4          D4
!! v           ^
!! ========    ========
!! v           ^
!! D4 > R4     D4 > R4
!! v           ^
!! D5 > R5     D5 > R5
!! v           ^
!! D6 > R6     D6 > R6
!! v           ^
!! D           D
!!
