module wei_fk_fx_wavefield_mod
  use wei_step_mod
  use wei_slow_mod
  implicit none


  contains
  

  subroutine store_wavefield(wfld,step,zero1)
    complex                 :: wfld(:,:,:,:,:)  !the wavefield
    type(westep)            :: step             !step object
    logical                 :: zero1            !wheter to zero tmp1 or tmp2
    integer                 :: i1,i2,i3,i4,iws


    iws=step%iw
  
    if(zero1) then
       do i4=1,size(wfld,4)
         do i3=1,size(wfld,3)
           do i2=1,size(wfld,2)
             do i1=1,size(wfld,1)
               step%tmp1(i1,i2,i3,i4)=0.
             end do
           end do
         end do
       end do
       do i4=1,size(wfld,4)
         do i3=1,size(wfld,3)
           do i2=1,size(wfld,2)
             do i1=1,size(wfld,1)
               step%tmp2(i1,i2,i3,i4)=wfld(i1,i2,i3,i4,iws)
             end do
           end do
         end do
       end do
     else
       do i4=1,size(wfld,4)
         do i3=1,size(wfld,3)
           do i2=1,size(wfld,2)
               do i1=1,size(wfld,1)
               step%tmp2(i1,i2,i3,i4)=0.
             end do
           end do
         end do
       end do
       do i4=1,size(wfld,4)
         do i3=1,size(wfld,3)
           do i2=1,size(wfld,2)
             do i1=1,size(wfld,1)
               step%tmp1(i1,i2,i3,i4)=wfld(i1,i2,i3,i4,iws)
             end do
           end do
         end do
       end do
    end if
  end subroutine


  !GRAB PART OF WAVEFIELD
  subroutine grab_part_wvfld(wfld,step)
    complex                 :: wfld(:,:,:,:,:)  !the wavefield
    type(westep)            :: step             !step object
    integer :: ith,iws,ifk,ierr,izs,iref,iomp,i3,i4,i,i3o,i4o
#ifdef DEBUG
    write(0,*) "IN GRAB PART WAVEFIELD"
#endif


    iws=step%iw
    do i4=1,size(wfld,4)
      do i3=1,size(wfld,3)
        call form_iref_slowness_slice(i3,i4,step,i3o,i4o)
        where(step%islow(:,:,i3o,i4o)==step%iref)
          wfld(:,:,i3,i4,iws) = step%tmp2(:,:,i3,i4)
        elsewhere
          wfld(:,:,i3,i4,iws) = 0.
        end where
      end do
    end do
#ifdef DEBUG
    write(0,*) "OUT GRAB PART WAVEFIELD"
#endif
  end subroutine

  !ADD WAVEFIELD
  subroutine add_wavefield(wfld,step)
    complex                 :: wfld(:,:,:,:,:)  !the wavefield
    type(westep)            :: step             !step object
    integer                 :: iws
    integer                 :: i1,i2,i3,i4
#ifdef DEBUG
    write(0,*) "IN ADD WAVEFIELD"
#endif
    iws=step%iw
     do i4=1,size(wfld,4)
       do i3=1,size(wfld,3)
         do i2=1,size(wfld,2)
           do i1=1,size(wfld,1)
             step%tmp1(i1,i2,i3,i4)=step%tmp1(i1,i2,i3,i4)+wfld(i1,i2,i3,i4,iws)
           end do
         end do
       end do
     end do
#ifdef DEBUG
     write(0,*) "OUT ADD WAVEFIELD"
#endif
  end subroutine

  !GRAB THE ENTIRE WAVEFIELD
  subroutine grab_wavefield(wfld,step,first)
    complex                 :: wfld(:,:,:,:,:)  !the wavefield
    type(westep)            :: step             !step object
    logical                 :: first            !whether or not store in first or second temp storage
    integer                 :: i1,i2,i3,i4
    integer                 :: iws
#ifdef DEBUG
    write(0,*) "IN GRAB WAVEFIELD"
#endif

    iws=step%iw
    if(first) then 
      do i4=1,size(wfld,4)
        do i3=1,size(wfld,3)
          do i2=1,size(wfld,2)
            do i1=1,size(wfld,1)
              wfld(i1,i2,i3,i4,iws) = step%tmp1(i1,i2,i3,i4)
            end do
          end do
        end do
      end do
    else
      do i4=1,size(wfld,4)
        do i3=1,size(wfld,3)
          do i2=1,size(wfld,2)
            do i1=1,size(wfld,1)
              wfld(i1,i2,i3,i4,iws) = step%tmp2(i1,i2,i3,i4)
            end do
          end do
        end do
      end do
    end if
#ifdef DEBUG
    write(0,*) "OUT GRAB WAVEFIELD"
#endif
  end subroutine

  !STORE PART OF THE WAVEFIELD INTO TEMPORARY BUFFER
  subroutine store_part_wavefield(wfld,step)
    complex                :: wfld(:,:,:,:,:)   !the wavefield
    type(westep)           :: step              !step object
    integer                :: i1,i2,i3,i4
    integer                :: iws,ierr,izs,i3o,i4o
#ifdef DEBUG
    write(0,*) "IN STORE PART WAVEFIELD"
#endif

    iws=step%iw
    do i4=1,size(wfld,4)
      do i3=1,size(wfld,3)
        write(0,*) "1XX",i3o,i4o,i3,i4
        call form_iref_slowness_slice(i3,i4,step,i3o,i4o)
        do i2=1,size(wfld,2)
          do i1=1,size(wfld,1)
            if(step%islow(i1,i2,i3o,i4o)==step%iref) then
              step%tmp2(i1,i2,i3,i4)=step%tmp2(i1,i2,i3,i4)+wfld(i1,i2,i3,i4,iws)
            end if
          end do
        end do
      end do
    end do
#ifdef DEBUG
    write(0,*) "IN STORE PART WAVEFIELD"
#endif
  end subroutine

end module
