module var_mod
  use external_mod
  use down_type_mod
  use run_mod
!  use lagrange_mod
  use var_type_mod
  implicit none
  integer,private,save :: npart,part(1000),iz(1000)
   integer,  private, pointer, save :: igood_len(:), do_freq_block(:)
  real, private, save,pointer :: zdepth(:,:),max_slow(:)

 contains

!-----------------------------------------------------
!
  logical function var_init()
#ifdef DEBUG
    real::a
    character(len=128), parameter :: unit='run_var_init'
    if(run%debug) call in(unit)
#endif DEBUG


    allocate(var%dzb(down%nzs))
    var_init=T

#ifdef DEBUG
    if(run%debug) call out(unit)
#endif DEBUG

  end function var_init



!-----------------------------------------------------
!
!
!   Initialize distribution pattern
!   (figure out what frequencies we are doing, where
!     they go, and what sampling we are using)
!
  logical function init_distrib_pattern()
!    type(dat), intent(in) :: X
    integer, pointer :: ifreq_list(:)
    real::a
    integer ::ierr
#ifdef DEBUG
    character(len=128), parameter :: unit='init_distrib_pattern'
    if(run%debug) call in(unit)
#endif

    init_distrib_pattern=F
    if(.not. create_frequency_list(ifreq_list)) then
       write(0,*) "trouble creating frequency list"
       return 
    end if

    allocate(down%iw_thread(size(ifreq_list)))

    if(down%order(1:3)/="VAR") then
       if(.not. distribute_reg(ifreq_list)) then
          write(0,*) "trouble distributing regular"
          return 
       end if
!    else 
!       if(.not. distribute_var_freq_slow(ifreq_list)) then
!          write(0,*) "trouble initializing varialbe freq slow pattern"
!          return
!       end if
    end if
    if(.not. store_distrib_pattern()) then
       write(0,*) "trouble storing distribution pattern"
       return
    end if
    init_distrib_pattern=T
#ifdef DEBUG
    if(run%debug) call out(unit)
#endif
  end function init_distrib_pattern
!-----------------------------------------------------
!  Create the list of frequencies we need to downward continue
logical function create_frequency_list(ifreq)
integer :: i,nredo,redo(10000)
integer, pointer :: ifreq(:)
#ifdef DEBUG
    character(len=128), parameter :: unit='create_frequency_list'
    if(run%debug) call in(unit)
#endif
create_frequency_list=F
! restart for old mpi version that divides by frequency
!if(.not. run%restart)  then 
  allocate(ifreq(down%aw%n/down%nws))
  do i=1,size(ifreq)
     ifreq(i)=i
  end do
!else 
!   nredo=getch("redo_groups","d",redo)
!   i=getch("redo_part","d",iz)
!   npart=getch("part_z","d",part)
!   if(i /= npart) &
!   write(0,*) "the number of redo_part and part_z not equal"
!   
!   allocate(ifreq(npart+nredo))
!   ifreq(1:nredo)=redo(1:nredo)
!   ifreq(nredo+1:nredo+npart)=part(1:npart)
!end if
down%nwp=size(ifreq)

create_frequency_list=T
#ifdef DEBUG
    if(run%debug) call out(unit)
#endif
  end function create_frequency_list
!-----------------------------------------------------
!
!
!  For the normal approach just distribute cycles
!
!
logical function distribute_reg(ifreq)
  integer, pointer :: ifreq(:)
  integer :: iw,isend

  distribute_reg=F


  if(mod(size(ifreq),run%m_nth)/=0) then
    if(run%verb) write(0,*) "WARNING--m_nth not a multiple of size(ifreq)",size(ifreq),run%m_nth
  end if

  down%nw_local=0;
  isend=0
  do iw=1,size(ifreq)
     down%iw_thread(iw)=isend
     if(run%debug) write(0,"(a,i3,a,i3)") "SENDING thread block",iw,&
     " to  processor ", down%iw_thread(iw)
     if(isend==run%m_ith) down%nw_local=down%nw_local+1
     isend=isend+1
     if(isend==run%m_nth) isend=0
  end do
!if(run%verb)write(0,*) 'list',down%iw_thread
 
 !Record the first frequency in each group
  allocate(down%iw_first(down%nw_local))
  isend=0
  do iw=1,size(ifreq)
     if(down%iw_thread(iw)==run%m_ith) then
        isend=isend+1
        down%iw_first(isend)=(ifreq(iw)-1)*down%nws+1
     end if
  end do
  distribute_reg=T
end function distribute_reg

!!--------------------------------------------------------------
!!  Distribute the frequencies with variable steps
!logical function distribute_var_freq_slow(ifreq)
!integer :: iw,i,k
!integer :: isend,nc,ierr,m,npts
!integer :: ifreq(:)
!real :: min_samp,max_samp,base_freq,io,lowcost
!integer :: cycles,nsamp,samples,itry,inode(1),igroup,j,nw,hxk
!real    ::  freq,lowest,myrand,mycost,min_zsamp,max_zsamp
!real,allocatable    :: mywork(:),dz(:,:),cost_node(:)
!integer, allocatable   :: nz(:)
!integer, allocatable   :: low_plan(:),cur_plan(:)
!integer :: nz_mem
!real :: mid
!logical,allocatable :: not_distrib(:)
!logical :: rite_sample
!#ifdef DEBUG
!    character(len=128), parameter :: unit='distribute_var_freq_slow'
!    if(run%debug) call in(unit)
!#endif
!distribute_var_freq_slow=F
!write(0,*) 'do i get into: distribute_var_freq_slow'
!call from_param("min_zsamp",min_zsamp,down%az%d/3.)
!call from_param("max_zsamp",max_zsamp,10000.)
!call from_param("cycles",cycles,1000000)
!call from_param("samp_per",samples,8)
!call from_param("io_cost",io,20.)
!isend=0
!down%nw_local=0
!
!nw=size(ifreq)
!allocate(dz(100000,nw))
!allocate(nz(nw))
!allocate(not_distrib(nw))
!allocate(low_plan(nw))
!allocate(cur_plan(nw))
!allocate(mywork(nw))
!allocate(cost_node(run%m_nth))
!
!!call return_max_slow(max_slow)
! max_slow=>down%max_slow
!if(run%verb) write(0,*) "CHEKC MAX SLOW",max_slow(1),max_slow(2),max_slow(3),max_slow(size(max_slow))
!
!do iw=1,nw
!  mid=down%nws*(ifreq(iw)-.5)
!  freq=mid*down%aw%d+down%aw%o
!  nz(iw)=calc_work(freq/6.285714,min_zsamp,max_zsamp,cycles,samples,dz(:,iw))
!  mywork(iw)=io+nz(iw)
!end do
!
!lowest=10000.
!
!!Find the plan that equalizes the work as much as possible
!do itry=1,10*nw
!  cost_node=0;
!
!  not_distrib=.true.
!
!  do while(any(not_distrib))
!    inode=minloc(cost_node)
!    call random_number(myrand)
!    igroup=nint(myrand*nw)
!   igroup=min(nw,max(1,igroup))
!    if(not_distrib(igroup)) then
!      cost_node(inode(1))=cost_node(inode(1))+mywork(igroup)
!      not_distrib(igroup)=.false.
!      cur_plan(igroup)=inode(1)
!    end if
!  end do
!
!  mycost=(maxval(cost_node)-minval(cost_node))/minval(cost_node)
!  if(mycost<lowest) then
!    lowest=mycost
!    low_plan=cur_plan
!  end if
!end do
!
!
!if(run%verb) write(run%io_unit,*) "cost differential",lowest
!
!low_plan=low_plan-1
!down%iw_thread=low_plan
!
!!Figure out how many groups local iw is holding
!down%nw_local=0
!do i=1,size(low_plan)
!  if(low_plan(i)==run%m_ith) down%nw_local=down%nw_local+1
!end do
!
!!Allocate the number of zs
!allocate(var%z_samp(down%nw_local))
!
!j=0
!allocate(down%iw_first(down%nw_local))
!
!   call from_param("nlagrange_pts",npts,5)
!
!
!do i=1,size(low_plan)
!  if(low_plan(i)==run%m_ith) then
!    j=j+1
!    down%iw_first(j)=(ifreq(i)-1)*down%nws+1
!    allocate(var%z_samp(j)%dz(nz(i)))
!    allocate(var%z_samp(j)%status(nz(i)))
!    allocate(var%z_samp(j)%zloc(nz(i)))
!    allocate(var%z_samp(j)%iloc(npts,down%az%n))
!    allocate(var%z_samp(j)%frac(npts,down%az%n))
!    var%z_samp(j)%dz(1:nz(i))=dz(1:nz(i),i)
!    var%z_samp(j)%zloc(1)=down%az%o
!     do m=2,nz(i)
!       var%z_samp(j)%zloc(m)=var%z_samp(j)%zloc(m-1)+var%z_samp(j)%dz(m-1)
!     end do
!   call lagrange_init(var%z_samp(j)%iloc,var%z_samp(j)%frac,var%z_samp(j)%zloc)
!  end if
!end do
!
!   call from_param("nz_mem",nz_mem,6)
!   call from_param("nlagrange_pts",npts,5)
!
!   allocate(var%filled(down%az%n))
!
!   !in core z buffer
!   if(nz_mem >0) then
!     allocate(var%v_i_buf(down%amx%n,down%amy%n,down%ahx%n,down%ahy%n,nz_mem))
!     allocate(var%mem_buf(nz_mem))
!     var%nmem_buf=nz_mem
!   else
!     var%nmem_buf=0
!     allocate(var%mem_buf(1))
!   end if
!
!
!   !on disk z buffer
!   if(npts-nz_mem >0) then
!     var%ndisk_buf=npts-nz_mem
!     var%use_disk_buf=.true.
!     allocate(var%disk_buf(npts-nz_mem))
!   else
!     var%NDISK_BUf=0
!     var%use_disk_buf=.false.
!     allocate(var%disk_buf(1))
!   end if
!
!distribute_var_freq_slow=T
!#ifdef DEBUG
!    if(run%debug) call out(unit)
!#endif
!end function distribute_var_freq_slow

!-------------------------------------------------------------
!
! Calculate the amount of work required for a given frequency

integer function calc_work(freq,min_dz,max_dz,max_cycle,sample_per,dz)
integer :: sample_per
logical :: slowness
integer :: n(3)
integer :: ipos,ithread,max_cycle
real :: d(3),freq,min_z,max_z,nax_dz,dv(3),max_Dz
real :: o(3),z,zstep,ov(3)
real :: dz(:)
integer i,npts,inew
real :: sample_hold(size(dz,1))
logical :: changed
integer :: iz
real :: fb,fe,min_dz,myf,cut,count
integer :: ib,ie,j,jb,je,ifreq,ierr,nv(3)
logical :: cutoff

nv=(/down%S%s_s%amx%n,down%S%s_s%amy%n,down%S%s_s%az%n/)
ov=(/down%S%s_s%amx%o,down%S%s_s%amy%o,down%S%s_s%az%o/)
dv=(/down%S%s_s%amx%d,down%S%s_s%amy%d,down%S%s_s%az%d/)
min_z=down%S%s_s%az%o
max_z=down%S%s_s%az%o +(down%S%s_s%az%n-1)*down%S%s_s%az%d


n=nv
d=dv
o=ov


myf=freq

z=min_z;npts=1;sample_hold(npts)=z
cutoff=.false.
count=0
do while(z< max_z .and. .not. cutoff)
  ipos=min(n(3),max(1,floor((z-o(3))/d(3)+1.5)))
  !GIVEN THE CURRENT VELOCITY DEPTH LOCATION, CHECK TO SEE HOW
  !LARGE IN DEPTH THIS FREQUENCY VALUE AND ARE SAMPLING
  zstep=(1./max_slow(ipos)/myf)/sample_per

  !check to see if there is a lower velocity with our given depth step
  inew=min(n(3),max(1,floor((z+zstep-o(3))/d(3)+1.5)))
  zstep=max(min_dz,(minval(1./max_slow(ipos:inew))/myf)/sample_per)
  zstep=min(max_dz,zstep)
  z=z+zstep
  npts=npts+1;sample_hold(npts)=z
  count=count + freq/maxval(max_slow(ipos:inew))* zstep
  if(count>max_cycle) cutoff=.true.
end do


!CALCULATE DZ SAMPLING
do i=1,npts-1
  dz(i)=sample_hold(i+1)-sample_hold(i)
end do

calc_work=npts


end function calc_work
!-------------------------------------------------------------
!
!  Find the next good fft length
!
!

integer function next_fast(nlow)
integer i,nlow
logical :: found

allocate(igood_len(230))
igood_len=(/ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,&
     18, 20, 21, 22, 24, 26, 28, 30, 33, 35, 36, 39, 40, 42, 44,&
     45, 48, 52, 55, 56, 60, 63, 65, 66, 70, 72, 77, 78, 80, 84,&
     88, 90, 91, 99, 104, 105, 110, 112, 117, 120, 126, 130, 132,&
    140, 143, 144, 154, 156, 165, 168, 176, 180, 182, 195, 198, 208,&
    210, 220, 231, 234, 240, 252, 260, 264, 273, 280, 286, 308, 312,&
    315, 330, 336, 360, 364, 385, 390, 396, 420, 429, 440, 455, 462,&
    468, 495, 504, 520, 528, 546, 560, 572, 585, 616, 624, 630, 660,&
    693, 715, 720, 728, 770, 780, 792, 819, 840, 858, 880, 910, 924,&
    936, 990, 1001, 1008, 1040, 1092, 1144, 1155, 1170, 1232, 1260,&
   1287, 1320, 1365, 1386, 1430, 1456, 1540, 1560, 1584, 1638, 1680,&
   1716, 1820, 1848, 1872, 1980, 2002, 2145, 2184, 2288, 2310, 2340,&
   2520, 2574, 2640, 2730, 2772, 2860, 3003, 3080, 3120, 3276, 3432,&
   3465, 3640, 3696, 3960, 4004, 4095, 4290, 4368, 4620, 4680, 5005,&
   5040, 5148, 5460, 5544, 5720, 6006, 6160, 6435, 6552, 6864, 6930,&
   7280, 7920, 8008, 8190, 8580, 9009, 9240, 9360, 10010, 10296,&
  10920, 11088, 11440, 12012, 12870, 13104, 13860, 15015, 16016,&
  16380, 17160, 18018, 18480, 20020, 20592, 21840, 24024, 25740,&
  27720, 30030, 32760, 34320, 36036, 40040, 45045, 48048, 51480,&
  55440, 60060, 65520, 72072, 80080, 90090, 102960, 120120, 144144,&
  180180, 240240, 360360, 720720/)

  i=1; found=.false.
  do while(i < size(igood_len) .and. .not. found)
    if(igood_len(i)< nlow) then
     i=i+1
    else
     found=.true.
    end if
  end do
  next_fast=igood_len(i)
end function
!-------------------------------------------------------------
!
!  Figure out the variable stepping for a given frequency
!
!


subroutine setup_var_block(iws,izs)
  integer :: iws,izs
  integer :: j,k,l,i,iz
  logical :: found
  if(down%order(1:3)/="VAR")  then
        var%dzb=1*down%az%d
  else
     j=(izs-1)*down%nzs
     k=min(size(var%z_samp(iws)%dz),down%nzs*izs)
     l=k-j
     var%dzb(1:l)=var%z_samp(iws)%dz(j+1:k)
     do iz=1,l
        i=1; found=.false.
        do while (i< size(zdepth,1) .and. .not. found )
           if(var%z_samp(iws)%zloc(iz+j) < zdepth(i,iws)) then
              found=.true.
           else
              i=i+1
           end if
        end do
     end do
  end if
end subroutine setup_var_block

!--------------------------------------------------------------
logical function  store_distrib_pattern()
  integer :: ierr
#ifdef DEBUG
    character(len=128), parameter :: unit='store_distrib_pattern'
    if(run%debug) call in(unit)
#endif
 store_distrib_pattern=F

if(run%m_ith==0) then
    call to_history("n1",size(down%iw_thread),"distrib_pattern")
    ierr=srite("distrib_pattern",real(down%iw_thread),size(down%iw_thread)*4)
end if
 store_distrib_pattern=T
#ifdef DEBUG
    if(run%debug) call out(unit)
#endif
end function store_distrib_pattern


!--------------------------------------------------------------
!
! Initialize the variable stepping buffers
!
!

  subroutine  var_samp_iw(iwb)
    integer :: iwb
    var%filled=.false.
    var%mem_buf=0
    var%disk_buf=0
    var%z_samp(iwb)%status=0
  end subroutine var_samp_iw

end module var_mod
module var_type_mod
  !---------------
  type  z_step
     real, pointer ::  dz(:)
     real, pointer ::  zloc(:)
     integer,pointer :: status(:)
     integer, pointer :: iloc(:,:)
     real, pointer :: frac(:,:)
  end type z_step
  !----------------------------------------------------------------
  ! A type to handle variable sampling
  type weivar
     character(len=128) :: operator
     character(len=128) :: operation,file_out
     type(z_step),pointer:: z_samp(:)   ! Sampling of local frequencies
     integer            :: distribute_pattern ! 0-reg,1-vel,2-vel+freq,
                                             ! 3-vel+freq+attenuation

     real,pointer       :: dzb(:) !z sampling for the current block
     integer,pointer    :: ihxb_beg(:) !beg loop for cur z block for hx axis
     integer,pointer    :: ihyb_beg(:) !beg loop for cur z block for hy axis
     integer,pointer    :: ihxb_end(:) !end loop for cur z block for hx axis
     integer,pointer    :: ihyb_end(:) !end loop for cur z block for hy axis
     complex,pointer    :: v_i_buf(:,:,:,:,:) !variable image buffer
     logical, pointer   :: filled(:) !whether or not output iz has been finished
     integer, pointer   :: disk_buf(:)!iz disk buffer element used
     integer, pointer   :: mem_buf(:)!iz memory buffer element used
     integer, pointer   :: write_freq(:) ! what block to start writing on
     integer            :: nmem_buf,ndisk_buf
     logical            :: use_disk_buf ! using disk buf
  end type weivar
  type(weivar) :: var
end module var_type_mod
