module wei_mem_mod
use sep
use wei_verb_mod
use sep3d_struct_mod
use wei_par_mod
use wei_run_mod
implicit none
integer, private, save :: mem_usage
logical, private, save :: verb_mem_usage,verb,my_reuse
integer, private, save :: ntot_fk,ntot_fx,ntot_wc,ntot_im
integer, private, save :: ibeg_fk,iend_fk
integer, private, save :: ibeg_fx,iend_fx
integer, private, save :: ibeg_wc,iend_wc
integer, private, save :: ibeg_im,iend_im
integer, private, save :: ibeg_rs,iend_rs
integer, private, save :: ibeg_wa,iend_wa,nomp
complex, allocatable, target :: wei_mem(:,:,:,:,:)
complex, pointer :: wei_mem2(:,:,:,:,:)

contains

  logical function init_mem_usage()

    call from_param('mem_monitor', verb_mem_usage, .false.)
    mem_usage=0
    ntot_wc=0
    ntot_im=0
    ntot_fk=0
    ntot_fx=0

    init_mem_usage=.true.
  end function

  subroutine grab_fx_limits(ib,ie)
   integer :: ib,ie
   ib=ibeg_fx
   ie=iend_fx
  end subroutine 

  subroutine grab_wc_limits(ib,ie)
   integer :: ib,ie
   ib=ibeg_wc
   ie=iend_wc
  end subroutine 

  subroutine grab_fk_limits(ib,ie)
   integer :: ib,ie
   ib=ibeg_fk
   ie=iend_fk
  end subroutine 

  subroutine set_fk_space(n)
   integer :: n
     ntot_fk=n
  end subroutine


  subroutine set_wc_space(n)
   integer :: n
     ntot_wc=n
  end subroutine


  subroutine set_fx_space(n)
   integer :: n
     ntot_wc=n
  end subroutine

  subroutine set_imaging_space(n)
   integer :: n
   ntot_im=n+ntot_im
  end subroutine

  subroutine assign_wave_space(array)
   complex, pointer :: array(:,:,:,:,:)
   array=>wei_mem(:,:,:,:,ibeg_wa:iend_wa)
  end subroutine

  subroutine assign_image_space(array1)
   complex, pointer          :: array1(:,:,:,:,:)
   array1=>wei_mem2(:,:,:,:,ibeg_rs:iend_rs)
  end subroutine

  subroutine assign_imaging_space(ith,array1,array2,array3)
   integer                   :: ith
   complex, pointer          :: array1(:,:,:,:,:)
   complex, pointer,optional :: array2(:,:,:,:,:)
   complex, pointer,optional :: array3(:,:,:,:,:)
   array1=>wei_mem(:,:,:,:,ibeg_im:ibeg_im+my_nomp()-1)
   if(present(array2)) array2=>wei_mem(:,:,:,:,ibeg_im+my_nomp():ibeg_im+2*my_nomp()-1)
   if(present(array3)) array3=>wei_mem(:,:,:,:,ibeg_im+my_nomp()*2:ibeg_im+3*my_nomp()-1)
  end subroutine

  subroutine assign_fx_space(ith,array1,array2,array3)
   integer                   :: ith
   complex, pointer          :: array1(:,:,:,:,:)
   complex, pointer,optional :: array2(:,:,:,:,:)
   complex, pointer,optional :: array3(:,:,:,:,:)
   array1=>wei_mem(:,:,:,:,ibeg_fx:ibeg_fx+my_nomp()-1)
   if(present(array2)) array2=>wei_mem(:,:,:,:,ibeg_fx+my_nomp():ibeg_fx+2*my_nomp()-1)
   if(present(array2)) array3=>wei_mem(:,:,:,:,ibeg_fx+my_nomp()*2:ibeg_fx+3*my_nomp()-1)
  end subroutine

  subroutine assign_wc_space(ith,array1,array2,array3)
   integer                   :: ith
   complex, pointer          :: array1(:,:,:,:,:)
   complex, pointer,optional :: array2(:,:,:,:,:)
   complex, pointer,optional :: array3(:,:,:,:,:)
   array1=>wei_mem(:,:,:,:,ibeg_wc:ibeg_wc+my_nomp()-1)
   if(present(array2)) array2=>wei_mem(:,:,:,:,ibeg_wc+my_nomp():ibeg_wc+2*my_nomp()-1)
   if(present(array2)) array3=>wei_mem(:,:,:,:,ibeg_wc+my_nomp()*2:ibeg_wc+3*my_nomp()-1)
  end subroutine

  subroutine assign_fk_space(ith,array1,array2,array3)
   integer                   :: ith
   complex, pointer          :: array1(:,:,:,:,:)
   complex, pointer,optional :: array2(:,:,:,:,:)
   complex, pointer,optional :: array3(:,:,:,:,:)
   array1=>wei_mem(:,:,:,:,ibeg_fk:ibeg_fk+my_nomp()-1)
   if(present(array2)) array2=>wei_mem(:,:,:,:,ibeg_fk+my_nomp():ibeg_fk+2*my_nomp()-1)
   if(present(array2)) array3=>wei_mem(:,:,:,:,ibeg_fk+my_nomp()*2:ibeg_fk+3*my_nomp()-1)
  end subroutine


 subroutine del_mem(iadd)
  integer  :: iadd
    mem_usage=mem_usage-iadd
 end subroutine

   subroutine show_mem(str)
    character(len=*) :: str
     if(my_Verb() .and. verb_mem_usage) &
      write(0,*) "memory used (mb) ",trim(str),":",real(mem_usage)*8./1000000.
   end subroutine

   subroutine add_mem(iadd)
    integer  :: iadd
      mem_usage=mem_usage+iadd
   end subroutine

   logical function reuse_mem()
     reuse_mem=my_reuse
   end function

   logical function wem_allocate_memory(wsep,rsep,n_omp)
    integer :: ierr,ntot,n_omp
    type(sep3d) :: wsep,rsep
    character(len=128) :: operation

    wem_allocate_memory=.false.

    nomp=n_omp
    call my_operation(operation)

    my_reuse=.false.
    if(wsep%n(3)==wsep%n(3) .and. wsep%d(3)==wsep%d(3) .and. &
        operation(1:6)=="migrat") then
      call from_param("reuse_image",my_reuse,.false.)
    end if
     

    !Doing one step at a time so we can reduce our memory usage
    if(wsep%n(7)==1 .and. wsep%n(8)==nomp) then
      ibeg_wa=1          ;   iend_wa=nomp             ; ntot=nomp  !wavefield
      if(my_reuse) then
        ntot=ntot+min(ntot_im*nomp+1,nomp*( ntot_fk+ntot_fx+ntot_wc))
        ibeg_rs=iend_wa+1  ;   iend_rs=iend_wa+1      
        ibeg_im=iend_rs+1  ;   iend_im=iend_rs+ntot_im*nomp

        ibeg_fk=iend_wa+1  ;   iend_fk=iend_wa+ntot_fk*nomp   
        ibeg_fx=iend_fk+1  ;   iend_fx=iend_fk+ntot_fx*nomp    
        ibeg_wc=iend_fx+1  ;   iend_wc=iend_fx+ntot_wc*nomp   
      else  
        ntot=ntot+min(ntot_im, ntot_fk+ntot_fx+ntot_wc)
        ibeg_im=iend_wa+1  ;   iend_im=iend_wa+ntot_im

        ibeg_fk=iend_wa+1  ;   iend_fk=iend_wa+ntot_fk*nomp   
        ibeg_fx=iend_fk+1  ;   iend_fx=iend_fk+ntot_fx*nomp    
        ibeg_wc=iend_fx+1  ;   iend_wc=iend_fx+ntot_wc*nomp   
        ibeg_rs=1          ;   iend_rs=1
      end if
      ibeg_im=iend_wa+1  ;   iend_im=iend_wa+ntot_im  ; ntot=ntot+ntot_im  !imaging oper
      
    else
      ibeg_wa=1          ;   iend_wa=wsep%n(8)                ; ntot=wsep%n(8)
      ibeg_fk=iend_wa+1  ;   iend_fk=iend_wa+ntot_fk*nomp     ; ntot=ntot+ntot_fk*nomp
      ibeg_fx=iend_fk+1  ;   iend_fx=iend_fk+ntot_fx*nomp     ; ntot=ntot+ntot_fx*nomp
      ibeg_wc=iend_fx+1  ;   iend_wc=iend_fx+ntot_wc*nomp     ; ntot=ntot+ntot_wc*nomp
      ibeg_im=iend_fx+1  ;   iend_im=iend_wc+ntot_im*nomp     ; ntot=ntot+ntot_im*nomp
      if(my_reuse) then
        ibeg_rs=iend_wc+1; iend_rs=iend_wc+wsep%n(7)        ; ntot=ntot+wsep%n(7)
      else 
        ibeg_rs=1        ; iend_rs=wsep%n(7)                ;
      end if 
    end if


       call show_mem("BEFORE BLOCK ALLOCATING")
       allocate(wei_mem(wsep%n(1),wsep%n(2),wsep%n(3),wsep%n(4),ntot),stat=ierr)
       if(ierr/=0) then
         write(0,*) "trouble allocating memory",sep_thread_num()
         return
       end if
       call add_mem(size(wei_mem))
       call show_mem("AFTER BLOCK ALLOCATING")
       if(.not. my_reuse) then
          allocate(wei_mem2(rsep%n(1),rsep%n(2),rsep%n(3),rsep%n(4),wsep%n(7)),stat=ierr)
          call add_mem(size(wei_mem2))
          call show_mem("AFTER IMAGE ALLOCATING")
       else
          wei_mem2=>wei_mem(:,:,:,:,ibeg_rs:iend_rs)
       end if 
       if(ierr/=0) then
         write(0,*) "trouble allocating memory",sep_thread_num()
         return
       end if
       wem_allocate_memory=.true.
   end function

end module
