module wei_z_mod
use wei_basics_mod
use wei_slow_mod
use wei_image_mod
use wei_data_mod
use sep3d_struct_mod
implicit none
type(sep3d), private, save              :: spath(2)
character(len=128) ,private,save        :: s_in(2),s_out(2)
integer ,private,save                   :: nslows
real, pointer    ,private,save          :: dzstep(:)
logical, pointer   ,private,save        :: iz_iw_do(:,:),iz_iw_io(:,:),iw_own(:)
logical        ,private,save            :: first_step,velocity
integer, pointer    ,private,save       :: i_z(:),i_w(:)


contains


!INITIALIZE EVERYTHING FOR SIMPLE DOWNWARD CONTINUATION IN Z
integer function path_z_init()

nullify(dzstep,iz_iw_do,iz_iw_io,i_z,i_w,iw_own)
call from_param("velocity",velocity,.false.)
path_z_init=.true.

end function



!
!SETUP THE SLOWNESS FILES USING THE Z COORDINATES
!
logical function  path_z_slow(wsep,rsep)
real :: first_z_step,first_z_vel
type(sep3d) :: ssep(2),wsep,rsep
real :: f
integer :: is,nsteps,istep,index(2),ifirst,ibeg,iend,ib,ie,i
real, allocatable :: slow(:,:),temp(:,:,:)
path_z_slow=.false.

!path is the same all the time so do nothing if already setup
  call from_param("first_z_step",first_z_step,0.)
  call from_param("first_z_vel",first_z_vel,1.5)
  nsteps=rsep%n(5)
  if(first_z_step>0.) then
    nsteps=nsteps+1
    first_step=.true.
  end if
  first_step=.false.
  if(associated(dzstep)) then
    call del_mem(size(dzstep))
    deallocate(dzstep)
  end if
  allocate(dzstep(nsteps))
  call add_mem(size(dzstep))

  !make a local slowness file with desired dz sampling
  call grab_slowness_masters(ssep,nslows,s_in)

  do is=1,nslows
    if(0/=sep3d_make_local_name("spath",s_out(is))) then
       write(0,*) "trouble making local slowness file name",sep_thread_num()
       return
    end if
    call init_sep3d(spath(is),"SCRATCH","FLOAT","REGULAR",(/1,1,1/))
    spath(is)%n(1:2)=ssep(is)%n(1:2)
    spath(is)%o(1:2)=ssep(is)%o(1:2)
    spath(is)%d(1:2)=ssep(is)%d(1:2)
    spath(is)%n(3)=nsteps
    call sep3d_set_sep3d(spath(is))
    call auxinout(s_out(is))
    call sep3d_write_description(s_out(is),spath(is))
    allocate(slow(spath(is)%n(1),spath(is)%n(2)))
    allocate(temp(spath(is)%n(1),spath(is)%n(2),2))
    call add_mem(size(slow)+size(temp))
    ifirst=0
    index=-1
    if(.not. sep3d_file_exist(s_in(is),.true.)) then
      write(0,*) "trouble opening",trim(s_out(is)),sep_thread_num()
      return
    end if
    if(first_z_step >0.) then
      ifirst=1
      dzstep(1)=first_z_step
      slow=1./first_z_vel
      if(.not. sep3d_write_data(s_out(is),spath(is),slow,local_tag=.true.,&
        fwind=(/0,0,0/),nwind=(/spath(is)%n(1),spath(is)%n(2),1/))) then  
          write(0,*) "trouble writing initial z step",sep_thread_num()
          return
      end if 
    end if
    do  istep=1,nsteps-ifirst
      f=wsep%o(5)+wsep%d(5)*(istep-1)  !location we want
      f=(f-ssep(is)%o(3))/ssep(is)%d(3)+1.     !location in terms of original model
      if(f< 1.) then
        ib=1; ie=2; f=0.
      else if(f>= ssep(is)%n(3)) then
        ib=ssep(is)%n(3)-1; ie=ssep(is)%n(3)-2; f=1.
      else
        ib=f; ie=ib+1
      end if
      if(ib/=index(1) .and. ib/=index(2)) then
        if(.not. update_slow_temp(s_in(is),ssep(is),index,ib,temp)) then
          write(0,*) "trouble updating slowness buffer",sep_thread_num()
          return
        end if
      end if
      if(ie/=index(1) .and. ie/=index(2)) then
        if(.not. update_slow_temp(s_in(is),ssep(is),index,ie,temp)) then
          write(0,*) "trouble updating slowness buffer",sep_thread_num()
          return
        end if
      end if
      f=f-floor(f)
      if(ib==index(1)) then
        slow=(1.-f)*temp(:,:,1)+f*temp(:,:,2)
      else
        slow=(1.-f)*temp(:,:,2)+f*temp(:,:,1)
      end if
      if(velocity) slow=1./slow
      if(.not. sep3d_write_data(s_out(is), spath(is),slow,local_tag=.true., &
         fwind=(/0,0,istep+ifirst-1/),  nwind=(/spath(is)%n(1),spath(is)%n(2),1/))) then
          write(0,*) "trouble writing local iz",istep
          return
      end if
    end do
    call del_mem(size(slow)+size(temp))
    deallocate(slow,temp)
    call auxclose(s_out(is))
 end do
 if(nslows==1) then
   call store_slowness(spath(1))
 else
   call store_slowness(spath(1),spath(2))
 end if
 dzstep(ifirst+1:)=wsep%d(5)
path_z_slow=.true.
end  function 

logical function update_slow_temp(tag,sep,index,iread,temp)
character(len=*)   :: tag
type(sep3d)        :: sep
integer            :: index(2),iread,ime,nh
real               :: temp(:,:,:)
update_slow_temp=.false.

ime=1
if(index(1)>index(2)) ime=2

call sep3d_grab_headers(tag,sep,nh,fwind=(/0,iread-1/),nwind=(/sep%n(2),1/),local_tag=.true.)

if(nh==0) write(0,*) "trouble reading space tag=",trim(tag)


if(.not. sep3d_read_data(tag,sep,temp(:,:,ime),local_tag=.true.)) then
  write(0,*) "trouble reading slowness section",iread," from tag",trim(tag)
  return
end if

index(ime)=iread

update_slow_temp=.true.
end function

logical function path_z_begin()
type(sep3d)      :: wsep,rsep,dsep,sep
character(len=128):: dtag,rtag
character(len=128)::  tag
integer          :: mode
real, pointer    :: dz(:),omega(:)
real, allocatable:: slow(:,:,:)
logical          :: adj
integer          :: i,i1,ido,i2,ndone,fdone

path_z_begin=.true.
call grab_current_wave_space(wsep)
call grab_current_image_space(rsep)
call grab_current_image_tag(rtag)
call grab_current_data_space(dsep)
call grab_current_data_tag(dtag)
call grab_current_adj(adj)

!FIRST CALL WITHIN MIGRATIOn
  if(.not.path_z_slow(wsep,rsep)) then
    write(0,*) "troublie initializing slowness"
    return
  end if
  if(associated(iz_iw_do)) then 
    call del_mem(size(iz_iw_do))
    deallocate(iz_iw_do)
  end if
  if(associated(iz_iw_io))then
    call del_mem(size(iz_iw_io))
    deallocate(iz_iw_io)
  end if
  if(associated(iw_own)) then
    call del_mem(size(iw_own))
    deallocate(iw_own)
  end if
  allocate(iz_iw_do(size(dzstep),dsep%n(5)))
  allocate(iz_iw_io(size(dzstep),dsep%n(5)))
  allocate(iw_own(wsep%n(6)))
  call add_mem(size(iz_iw_do)+size(iz_iw_io)+size(iw_own))
  do i=1,wsep%n(6)
    iw_own(i)=sep3d_own(dsep,i)
  end do
  sep=rsep;if(.not. adj) sep=dsep
  if(.not. get_work(adj,wsep,sep,iz_iw_do,iz_iw_io)) then
    write(0,*) "trouble grabbing work",sep_thread_num()
    return
  end if 
  fdone=0; ndone=0;
  do i=1,wsep%n(6)
    if(iw_own(i)) then
      ndone=ndone+1
      if(all(iz_iw_do(:,i))) fdone=fdone+1
    end if
  end do
  if(.not. init_print_progress(fdone*size(iz_iw_do,1),ndone*size(iz_iw_do,1))) then
    write(0,*) "trouble intitializing verbosing"
    return
  end if


path_z_begin=.true.
end function






!
! Go along the Z axis
!
logical function path_z_op(mode,data,image,dz,omega)
type(sep3d)      :: wsep,rsep,dsep,sep
character(len=128):: dtag,rtag
character(len=128)::  tag
integer          :: mode
complex          :: data(:,:,:,:,:)
complex             :: image(:,:,:,:,:)
real, pointer    :: dz(:),omega(:)
real, allocatable:: slow(:,:,:)
logical          :: adj
integer          :: i,i1,ido,i2,ndone,fdone

path_z_op=.false.
call grab_current_wave_space(wsep)
call grab_current_image_space(rsep)
call grab_current_image_tag(rtag)
call grab_current_data_space(dsep)
call grab_current_data_tag(dtag)
call grab_current_adj(adj)

!FIRST CALL WITHIN MIGRATIOn

if(mode==mode_begin) then
  if(associated(omega)) then
     call del_mem(size(omega))
     deallocate(omega)
     nullify(omega)
  end if
  if(associated(dz)) then
     call del_mem(size(dz))
     deallocate(dz)
     nullify(dz)
  end if
  !Figure out what we need to do
  call find_iw(adj,dsep,rsep,wsep,iz_iw_do,i_w)          !what omega haven't been finished
  if(.not. associated(i_w)) then
    path_z_op=1; return
  end if
  call find_iz(adj,wsep,iz_iw_do,i_w,i_z)!find what dz we need to
  if(.not. read_slow_block(i_z)) then          !read in a slowness blcok
    write(0,*) "trouble reading in slowness block"
    return
  end if
  if(adj) then
    if(i_z(1)==1) then 
       if(.not. read_data_block(dtag,dsep,i_w,data)) then
         write(0,*) "trouble2reading data block"
         return
       end if
    end if
    if(.not. reuse_mem()) image=0.; !zero the image block
  else
    if(.not. read_image_block(rtag,rsep,i_z,image)) then
      write(0,*) "trouble reading image block"
      return 
    end if
  end if
  allocate(omega(size(i_w))); omega=wsep%o(6)+wsep%d(6)*(i_w-1)
  allocate(dz(size(i_z))); 
  call add_mem(size(dz)+size(omega))
  do i=1,size(i_z)
    dz(i)=dzstep(i_z(i))
  end do
else if(mode==mode_end) then
  if(adj) then
    if(.not. write_image_block(rtag,rsep,i_z,image)) then
      write(0,*) "trouble writing image block"
      return
    end if
  else
    if(i_z(size(i_z))==1) then 
      if(.not. write_data_block(dtag,dsep,i_w,data)) then
        write(0,*) "trouble writing data block"
        return
      end if
    end if
  end if
  do i2=1, size(i_w)
    do i1=1,size(i_z)
      iz_iw_do(i_z(i1),i_w(i2))=.true.
      iz_iw_io(i_z(i1),i_w(i2))=.true.
    end do
  end do
else 
  write(0,*) "invalid mode",mode
end if

path_z_op=.true.
end function

!
!  Figure out what frequencies and depths we need to do
!
logical function get_work(adj,wsep,sep,iz_iw_do,iz_iw_io)
logical     :: adj
type(sep3d) :: sep,wsep 
character(len=55) ::  par
integer, allocatable :: iw(:),iz(:)
integer     :: ifreq,ii,i,i2,i1
logical     :: iz_iw_do(:,:),iz_iw_io(:,:)
get_work=.false.
iz_iw_do=.false.
iz_iw_io=.false.
if(my_restart()) then
  allocate(iz(sep_num_thread()),iw(sep_num_thread()))
  call add_mem(size(iz)+size(iw))
 if(.not. sep3d_grab_sect_param(sep,"iz_done",iz)) then
    write(0,*) "trouble grabbing section parameters"
    return
  end if
  do ifreq=1,wsep%n(6) !loop over frequencies
    par="ifreq_done"
    call parcat(par,ifreq)
    iw=-1;
    if(.not. sep3d_grab_sect_param(sep,par,iw)) then
      write(0,*) "trouble grabbing section parameters"
      return
    end if
    if(any(iw==1)) then
      iz_iw_io(:,ifreq)=.true.
      iz_iw_do(:,ifreq)=.true.
    else if(any(iw==0)) then
      do i=1,sep_num_thread()
        if(iw(i)==0) ii=i
      end do
      if(adj) iz_iw_io(1:iz(ii),ifreq)=.true.
    end if
  end do
  call del_mem(size(iz)+size(iw))
  deallocate(iw,iz)
end if
if(first_step) iz_iw_io(1,:)=.true. !Don't write out the water step
get_work=.true.
end function

!
!FIND THE FIRST IW TO WORK that hasn't been computer
!
subroutine find_iw(adj, dsep, rsep,wsep,iz_iw_do,i_w)
logical            :: adj
type(sep3d)        :: wsep,dsep,rsep
logical            :: iz_iw_do(:,:)
integer,pointer            :: i_w(:)
integer            :: n,i
character(len=64)  :: par
n=0;

if(associated(i_w))then
  if(any(.not. iz_iw_do(:,i_w(1)))) return
  do i=1,size(i_w) 
     par="ifreq_done"
     call parcat(par,i_w(i))
    if(adj) then
      if(.not. sep3d_set_sect_param(rsep,par,1)) then
        write(0,*) "trouble writing sect_params"
        return
      end if
    else
      if(.not. sep3d_set_sect_param(dsep,par,1)) then
        write(0,*) "trouble writing sect_params"
        return
      end if
    end if
  end do
  call del_mem(i_w)
  deallocate(i_w)
end if
nullify(i_w)
do i=1,wsep%n(6)
 if(iw_own(i) .and.  any(.not.iz_iw_do(:,i))) n=n+1
end do
n=min(n,wsep%n(8))
if(n>0) then
  allocate(i_w(n))
  call add_mem(size(i_w))
  n=0;
  do i=1,wsep%n(6)
    if(iw_own(i) .and.  any(.not.iz_iw_do(:,i)).and. n < size(i_w))then
      n=n+1
      i_w(n)=i
    end if
  end do
  do i=1,size(i_w) 
    par="ifreq_done"
    call parcat(par,i_w(i))
    if(adj) then
      if(.not. sep3d_set_sect_param(rsep,par,0)) then
        write(0,*) "trouble writing sect_param"
        return
      end if
    else
      if(.not. sep3d_set_sect_param(dsep,par,0)) then
        write(0,*) "trouble writing sect_param"
        return
      end if
    end if
  end do
end if
end subroutine

!FIND THE FIRST IZ TO WORK ON
subroutine find_iz(adj,wsep,iz_iw_do,i_w,i_z)
logical            :: adj
type(sep3d)        :: wsep
logical,pointer            :: iz_iw_do(:,:)
integer,pointer            :: i_w(:),i_z(:)
integer            :: n,i,ibeg,iend,istep,ifirst
logical, pointer   :: pt_it(:,:)


call grab_current_adj(adj)
if(.not. associated(i_w)) return !we have no frequencies
if( associated(i_z)) then
   call del_mem(size(i_z))
   deallocate(i_z)
end if

allocate(pt_it(size(iz_iw_do,1),size(i_w)))
call add_mem(size(pt_it))
do i=1,size(i_w)
  pt_it(:,i)=iz_iw_do(:,i_w(i))
end do


ibeg=1; iend=size(iz_iw_do,1); istep=1
if(.not. adj)  then
  iend=1; ibeg=size(iz_iw_do,1); istep=-1
end if
i=ibeg;
ifirst=-1;
do while(i <=iend .and. ifirst==-1) 
  if(all (.not. pt_it(i,:))) ifirst=i

  i=i+istep
end do
if(adj) then
  n=min(size(iz_iw_do,1)-ifirst+1,wsep%n(7))
  allocate(i_z(n))
  call add_mem(size(i_z))
  do i=1,n
    i_z(i)=ifirst-1+i
  end do
else
  n=min(ifirst,wsep%n(7))
  allocate(i_z(n))
  call add_mem(size(i_z))
  do i=1,n
    i_z(i)=ifirst+1-i
  end do
end if
call del_mem(size(pt_it))
deallocate(pt_it)
end subroutine


!
!Read in slowness block
!
logical function read_slow_block(i_z)
integer :: i_z(:),nh
real, pointer :: slow1(:,:,:),slow2(:,:,:)
read_slow_block=.false.

if(nslows==1) then
   call grab_slowness_block(slow1)
else
   call grab_slowness_block(slow1,slow2)
end if


call sep3d_grab_headers(s_out(1),spath(1),nh,fwind=(/0,minval(i_z)-1/),&
   nwind=(/spath(1)%n(2),size(i_z)/),local_tag=.true.)
if(.not. sep3d_read_data(s_out(1),spath(1),slow1,local_tag=.true.)) then
  write(0,*) "trouble reading slowness"
  return
end if
if(nslows>1) then
  call sep3d_grab_headers(s_out(2),spath(2),nh,fwind=(/0,minval(i_z)-1/),&
     nwind=(/spath(2)%n(2),size(i_z)/),local_tag=.true.)
  if(.not. sep3d_read_data(s_out(2),spath(1),slow2,local_tag=.true.)) then
    write(0,*) "trouble reading slowness"
    return
  end if
end if

read_slow_block=.true.
end  function

!
!Write out data block
!
logical function write_data_block(tag,sep,iax,array)
character(len=*)       :: tag
type(sep3d)            :: sep
integer                :: iax(:)
complex                :: array(:,:,:,:,:)
integer                :: i
character(len=64)      :: par
write_data_block=.false.

do i=1,size(iax)
  if(.not.write_data_if(tag,sep,iax(i),i,array)) then
    write(0,*) "trouble writing data if"
    return
  end if
  par="ifreq_done"
  call parcat(par,iax(i))
  if(.not. sep3d_set_sect_param(sep,par,1)) then
    write(0,*) "trouble writing sect_param"
    return
  end if
end do

write_data_block=.true.
end function


!
!Write out image block
!
logical function write_image_block(tag,sep,iax,array)
character(len=*)       :: tag
type(sep3d)            :: sep
integer                :: iax(:),i
complex                :: array(:,:,:,:,:)
write_image_block=.false.

do i=1,size(iax)
  if(.not.write_image_iz(tag,sep,iax(i),i,array)) then
    write(0,*) "trouble writing image block"
    return
  end if
  if(.not. sep3d_set_sect_param(sep,"iz_done",iax(i))) then
    write(0,*) "trouble writing sect_param"
    return
  end if
end do

write_image_block=.true.
end function

!
!Read in data block
!

logical function read_data_block(tag,sep,iax,array)
character(len=*)       :: tag
type(sep3d)            :: sep
integer                :: iax(:),i
complex                :: array(:,:,:,:,:)
read_data_block=.false.
do i=1,size(iax)
  if(.not.read_data_if(tag,sep,iax(i),i,array)) then
    write(0,*) "trouble1reading data block",iax(i)
    return
  end if
end do
read_data_block=.true.

end function


!
!Read in image block
!
logical function read_image_block(tag,sep,iax,array)
character(len=*)       :: tag
type(sep3d)            :: sep
integer                :: iax(:),i
complex                :: array(:,:,:,:,:)
read_image_block=.false.
do i=1,size(iax)
  if(.not.read_image_iz(tag,sep,iax(i),i,array)) then
    write(0,*) "trouble reading image block"
    return
  end if
end do
read_image_block=.false.
end function

end module
