module wei_wem_ssf_scale_mod
  use wei_ssf_mod
  use wei_wem_mod
  use wei_slow_mod
  use wei_dsr_mod
  use wei_damp_mod
  use wei_slow_mod
  use wei_fk_fx_wavefield_mod
  use sep_timers_mod
  use wei_fft_mod
  implicit none
  integer,private, save :: tm1,tm2,tm3,tm4,tm5,tm6,tm7
  real,private, save :: sfact



  contains

 subroutine wem_ssf_single()

   call from_param("sfact",sfact,1000000.)

 end subroutine


  subroutine wem_ssf_init(wfld,step)
    type(westep)         :: step
    complex              :: wfld(:,:,:,:,:)
    logical              :: logit

#ifdef DEBUG
    write(0,*) "IN WEM SSF INIT"
#endif

if(step%ithread==1)    logit=setup_next_timer("OTHER1STEP",tm1)
if(step%ithread==1)     logit=setup_next_timer("OTHER2STEP",tm5)
    if(step%ithread==1) logit=setup_next_timer("OTHER3STEP",tm6)
    if(step%ithread==1) logit=setup_next_timer("OTHER4STEP",tm7)
    if(step%ithread==1) logit=setup_next_timer("WEI STEP",tm2)
    if(step%ithread==1) logit=setup_next_timer("SSF STEP",tm3)
    if(step%ithread==1) logit=setup_next_timer("FFT STEP",tm4)

    call wem_init(step%wave,step%fk)   !Initialzie WEM operation
    call wei_fft_init_par(step)       !Initialize FFT
    call build_padding(step,(step%wave%npad(1:4)-step%wave%n(1:4))) !BUILD PADDING

    !Initialize temporary storage for wavefields
    allocate(step%tmp1(step%wave%npad(1),step%wave%npad(2),step%wave%npad(3),step%wave%npad(4)))
    allocate(step%tmp2(step%wave%npad(1),step%wave%npad(2),step%wave%npad(3),step%wave%npad(4)))
    
   

#ifdef DEBUG
    write(0,*) "OUT WEM SSF INIT"
#endif

  end subroutine



  subroutine one_step(wfld,step,wave)
    complex      :: wfld(:,:,:,:,:)   !wavefield
    type(westep) :: step
    type(wspace) :: wave
    type(ctable) :: dsr
#ifdef DEBUG
     write(0,*) "IN ONE STEP"
#endif


    call grab_dsr(step,dsr)

    if(step%wave%adj) then
      call set_source(step,step%wave)
      call one_step_down(wfld,step,dsr)
      call set_receiver(step)
      call one_step_down(wfld,step,dsr)
    else 
      call set_receiver(step)
      call one_step_up(wfld,step,dsr)
      call set_source(step,step%wave)
      call one_step_up(wfld,step,dsr)
    end if    
#ifdef DEBUG
     write(0,*) "OUT ONE STEP"
#endif
  end subroutine

  subroutine one_step_up(wfld,step,dsr) !s g migration for example
    complex      :: wfld(:,:,:,:,:)   !wavefield
    type(westep) :: step              !current step
    type(ctable) :: dsr               !dsr table
    integer      :: iref              !current reference velocity
#ifdef DEBUG
     write(0,*) "IN ONE STEP UP"
#endif

    call damp_wave(wfld,step) !damp the stored wavefield
    call wei_ffft(wfld,step)  !transform into fk
    if(size(step%srefs)==1) then
      call set_reference(step,1)
      call wei_wem(wfld,step,dsr)
      call wei_ifft(wfld,step)
      call wei_ssf(wfld,step)
    else
      call store_wavefield(wfld,step,.true.) !STORE THE CURRENT WAVEFIELD
      do iref=1,size(step%srefs)
        call set_reference(step,iref)
        call grab_wavefield(wfld,step,.true.)
        call wei_wem(wfld,step,dsr)
        call wei_ifft(wfld,step)
        call wei_ssf(wfld,step)
        call store_part_wavefield(wfld,step)  !add fx
      end do
      call grab_wavefield(wfld,step,.false.);
    end if
#ifdef DEBUG
     write(0,*) "OUT ONE STEP UP"
#endif
  end subroutine


  subroutine one_step_down(wfld,step,dsr) !s g migration for example
    complex      :: wfld(:,:,:,:,:)   !wavefield
    type(westep) :: step              !current step
    type(ctable) :: dsr               !dsr table
    integer      :: iref              !current reference velocity

#ifdef DEBUG
     write(0,*) "IN ONE STEP DOWN"
#endif
    if(size(step%srefs)==1) then
      call set_reference(step,1)
      call wei_ssf(wfld,step)
        call scale(wfld,step%iw)
      call wei_ifft(wfld,step)
        call scale(wfld,step%iw)
      call wei_wem(wfld,step,dsr)
        call scale(wfld,step%iw)
    else
if(step%ithread==1) call start_timer_num(tm1)
      call store_wavefield(wfld,step,.true.) !STORE THE CURRENT WAVEFIELD
if(step%ithread==1)call stop_timer_num(tm1)
      do iref=1,size(step%srefs)
if(step%ithread==1)call start_timer_num(tm5)
        call set_reference(step,iref)
        call grab_part_wvfld(wfld,step)
if(step%ithread==1)call stop_timer_num(tm5)
if(step%ithread==1)call start_timer_num(tm3)
        call wei_ssf(wfld,step)
        call scale(wfld,step%iw)
if(step%ithread==1)call stop_timer_num(tm3)
if(step%ithread==1)call start_timer_num(tm4)
        call wei_ifft(wfld,step)
        call scale(wfld,step%iw)
if(step%ithread==1)call stop_timer_num(tm4)
if(step%ithread==1)call start_timer_num(tm2)
        call wei_wem(wfld,step,dsr)
        call scale(wfld,step%iw)
if(step%ithread==1)call stop_timer_num(tm2)
if(step%ithread==1)call start_timer_num(tm6)
        call add_wavefield(wfld,step)
if(step%ithread==1)call stop_timer_num(tm6)
      end do
if(step%ithread==1)call start_timer_num(tm7)
      call grab_wavefield(wfld,step,.true.)
if(step%ithread==1)call stop_timer_num(tm7)
    end if
if(step%ithread==1)call start_timer_num(tm4)
    call wei_ffft(wfld,step)
if(step%ithread==1)call stop_timer_num(tm4)
if(step%ithread==1)call start_timer_num(tm7)
!    call damp_wave(wfld,step)
if(step%ithread==1)call stop_timer_num(tm7)
#ifdef DEBUG
     write(0,*) "OUT ONE STEP DOWN"
#endif
  end subroutine
  subroutine scale(wfld,iw)
    complex :: wfld(:,:,:,:,:)
    integer :: iw
    integer :: i1,i2,i3,i4
    real    :: r,i

    do  i4=1,size(wfld,4)
    do  i3=1,size(wfld,3)
    do  i2=1,size(wfld,2)
    do  i1=1,size(wfld,1)
       r=nint(real(wfld(i1,i2,i3,i4,iw)*sfact))/sfact
       i=nint(aimag(wfld(i1,i2,i3,i4,iw)*sfact))/sfact
       wfld(i1,i2,i3,i4,iw)=cmplx(r,i)
    end do
    end do
    end do
    end do
  end subroutine

end module
