module wei_image_mod
use wei_wavespace_mod
implicit none
type(sep3d),private,save :: rsep
character(len=128),private,save :: rtag


contains

!-------------------------------------------------!
!                                                 !
!         INITIALIZATION ROUTINES                 !
!                                                 !
!-------------------------------------------------!
logical function init_image(tag,myrsep,created)
type(sep3d) :: myrsep
character(len=*) :: tag
logical          :: created

init_image=.false.

if(sep3d_file_exist(tag)) then
  if(.not. init_image_tag(tag,myrsep)) then
    write(0,*) "trouble initializing tag"
    return
  end if
  created=.false.
else
  if(.not. init_image_struct(tag,myrsep)) then
    write(0,*) "trouble initializing structg"
    return
  end if
  created=.true.
end if
init_image=.true.

end function



logical function init_image_struct(tag,myrsep)
character(len=*)  :: tag
type(sep3d) :: myrsep,wsep
logical :: myre

init_image_struct=.false.

call  grab_current_wave_space(wsep)
if(.not. image_distrib_tag(wsep,myrsep,.false.)) then
  write(0,*) "trouble initializing distribution"
  return
end if

if(sep_thread_num()==0) call auxinout(tag)
call sep3d_write_description(tag,myrsep)
if(sep_thread_num()==0) then
  if(.not. zero_image(tag,myrsep)) then
    write(0,*) "trouble zeroing image"
    return
   end if
end if

init_image_struct=.true.
end function


logical function image_distrib_tag(wsep,myrsep,spread)
type(sep3d) :: wsep,myrsep
logical :: myre,spread
image_distrib_tag=.false.

call from_param("image_real",myre,.false.)
if(myre) then
  call init_sep3d(myrsep,"SCRATCH","FLOAT","REGULAR",(/1,1,1,1,1,1/))
else
  call init_sep3d(myrsep,"SCRATCH","COMPLEX","REGULAR",(/1,1,1,1,1,1/))
end if
myrsep%n(1:5)=wsep%n(1:5)
myrsep%o(1:5)=wsep%o(1:5)
myrsep%d(1:5)=wsep%d(1:5)
myrsep%label(1)="CMP_X"
myrsep%label(2)="CMP_Y"
myrsep%label(3)="H_X"
myrsep%label(4)="H_Y"
myrsep%label(5)="DEPTH"
call sep3d_set_sep3d(myrsep)

if(spread .and. sep_num_thread()>1) then
  myrsep%n(6)=sep_num_thread()
  call sep3d_set_sep3d(myrsep)
  if(.not. sep3d_section_tag(myrsep,6,"BLOCK",sep_num_thread())) then
    write(0,*) "trouble sectioning tag"
    return
  end if
end if


image_distrib_tag=.true.
end function

logical function init_image_tag(tag,myrsep)
character(len=*) :: tag
type(sep3d) :: myrsep
init_image_tag=.false.

call init_sep3d(tag,myrsep,"SCRATCH")
init_image_tag=.true.

end function

!-------------------------------------------------!
!                                                 !
!                   IO ROUTINES                   !
!                                                 !
!-------------------------------------------------!

logical function write_image_iz(tag,myrsep,iz,ielem,wfld)
type(sep3d) :: myrsep
character(len=*) :: tag
complex :: wfld(:,:,:,:,:)
integer ::i3,i4,i5,nw(6),fw(6),ndim,nalloc,nloop,nh,iz,ielem
logical :: ddd,big
complex, allocatable :: junkc(:,:)
real, allocatable :: junkr(:,:)
write_image_iz=.false.
 

big=.false.

nw(1:5)=myrsep%n(1:5); nw(4)=1; nw(5)=1; nw(3)=1;nw(6)=1
fw=0; fw(6)=sep_thread_num()

if(myrsep%n(1)*myrsep%n(2)< 1000000) then
   big=.true.
   nalloc=myrsep%n(2)*myrsep%n(3)
   nw(3)=myrsep%n(3)
   nloop=1
else
   nalloc=myrsep%n(2)
   nloop=myrsep%n(3)
end if


if(myrsep%data_format(1:1)=="C") then
  if(big) then
    allocate(junkc(myrsep%n(1),nalloc),junkr(1,1))
  else
    allocate(junkc(1,1),junkr(1,1))
  end if 
else
  allocate(junkr(myrsep%n(1),nalloc),junkc(1,1))
end if
call add_mem(size(junkc)+size(junkr))

fw(5)=iz-1
    do i4=1,myrsep%n(4)
      fw(4)=i4-1
      do i3=1,nloop
        fw(3)=i3-1
        call sep3d_grab_headers(tag,myrsep,nh,nwind=nw(2:),fwind=fw(2:))
        if(myrsep%data_format(1:1)=="C") then
          if(big) then
            if(.not. sep3d_read_data(tag,myrsep,junkc)) then
              write(0,*) "trouble reading image",iz
              return
            end if
            junkc=reshape(wfld(:,:,:,i4,ielem),(/myrsep%n(1),nalloc/))+junkc
            ddd=sep3d_write_data(tag,myrsep,junkc,nwind=nw,fwind=fw)
          else 
            ddd=sep3d_write_data(tag,myrsep,wfld(:,:,i3,i4,ielem),nwind=nw,fwind=fw)
          end if
        else 
          if(.not. sep3d_read_data(tag,myrsep,junkr)) then
            write(0,*) "trouble reading image" ,iz
            return
          end if
          if(big) then
            junkr=reshape(real(wfld(:,:,:,i4,ielem)),(/myrsep%n(1),nalloc/))+junkr
          else
            junkr=reshape(real(wfld(:,:,i3,i4,ielem)),(/myrsep%n(1),nalloc/))+junkr
          end if
          ddd=sep3d_write_data(tag,myrsep,junkr,nwind=nw,fwind=fw)
        end if
        if(.not. ddd) then
          call del_mem(size(junkc)+size(junkr))
          deallocate(junkr,junkc)
          write(0,*) "trouble writing iz",i3,i4,iz,sep_thread_num()
          return
        end if
      end do
    end do
call del_mem(size(junkc)+size(junkr))
deallocate(junkr,junkc)
write_image_iz=.true.
end function

logical function read_image_iz(tag,myrsep,iz,ielem,wfld)
type(sep3d) :: myrsep
character(len=*) :: tag
complex :: wfld(:,:,:,:,:)
complex, allocatable :: junkc(:,:)
real, allocatable :: junkr(:,:)
integer ::i3,i4,i5,nw(6),fw(6),ndim,nalloc,nloop,nh,iz,ielem
logical :: ddd,big
read_image_iz=.false.
 

big=.false.

nw(1:5)=myrsep%n(1:5); nw(4)=1; nw(5)=1; nw(3)=1; nw(6)=1 ;nw(5)=1
fw=0
fw(5)=iz-1;fw(6)=sep_thread_num()

if(myrsep%n(1)*myrsep%n(2)< 1000000) then
   big=.true.
   nalloc=myrsep%n(2)*myrsep%n(3)
   nw(3)=myrsep%n(3)
   nloop=1
else
   nalloc=myrsep%n(2)
   nloop=myrsep%n(3)
end if


if(myrsep%data_format(1:1)=="C") then
  if(big) then
    allocate(junkc(myrsep%n(1),nalloc),junkr(1,1))
  else
    allocate(junkc(1,1),junkr(1,1))
  end if 
else
  allocate(junkr(myrsep%n(1),nalloc),junkc(1,1))
end if
call add_mem(size(junkr)+size(junkc))

    do i4=1,myrsep%n(4)
      fw(4)=i4-1
      do i3=1,nloop
        fw(3)=i3-1
        call sep3d_grab_headers(tag,myrsep,nh,nwind=nw(2:),fwind=fw(2:))
        if(myrsep%data_format(1:1)=="C") then
          if(big) then
            ddd=sep3d_read_data(tag,myrsep,junkc)
            wfld(:,:,:,i4,ielem)=reshape(junkc,(/myrsep%n(1),myrsep%n(2),myrsep%n(3)/))
          else 
            ddd=sep3d_read_data(tag,myrsep,wfld(:,:,i3,i4,ielem))
          end if
        else 
          ddd=sep3d_read_data(tag,myrsep,junkr)
          if(big) then
            wfld(:,:,:,i4,ielem)=cmplx(reshape(junkr,(/myrsep%n(1),myrsep%n(2),myrsep%n(3)/)),0.)
          else
            wfld(:,:,i3,i4,ielem)=cmplx(reshape(junkr,(/myrsep%n(1),myrsep%n(2)/)),0.)
          end if
        end if
        if(.not. ddd) then
          call del_mem(size(junkr)+size(junkc))
          deallocate(junkr,junkc)
          write(0,*) "trouble writing iz",i3,i4,iz,sep_thread_num()
          return
        end if
      end do
    end do
call del_mem(size(junkr)+size(junkc))
deallocate(junkr,junkc)
read_image_iz=.false.
end function


logical function zero_image(tag,myrsep)
type(sep3d) :: myrsep
character(len=*) :: tag
real, allocatable :: junkr(:,:)
complex, allocatable :: junkc(:,:)
integer ::i3,i4,i5,nw(6),fw(6),ndim
logical :: ddd
zero_image=.false.
if(myrsep%data_format(1:1)=="C") then
  allocate(junkc(myrsep%n(1),myrsep%n(2)),junkr(1,1))
  junkc=0
else
  allocate(junkr(myrsep%n(1),myrsep%n(2)),junkc(1,1))
  junkr=0
end if
call add_mem(size(junkc)+size(junkr))

nw(1:5)=myrsep%n(1:5); nw(4)=1; nw(5)=1; nw(3)=1
fw=0;fw(6)=sep_thread_num(); nw(6)=1; nw(5)=1

do i5=1,myrsep%n(5)
  fw(5)=i5-1
    do i4=1,myrsep%n(4)
      fw(4)=i4-1
      do i3=1,myrsep%n(3)
        fw(3)=i3-1
        if(myrsep%data_format(1:1)=="C") then
          ddd=sep3d_write_data(tag,myrsep,junkc,nwind=nw,fwind=fw)
        else 
          ddd=sep3d_write_data(tag,myrsep,junkr,nwind=nw,fwind=fw)
        end if
        if(.not. ddd) then
          deallocate(junkr,junkc)
          write(0,*) "trouble zero part",i3,i4,i5,sep_thread_num()
          return
        end if
      end do
    end do
end do
call del_mem(size(junkc)+size(junkr))
deallocate(junkr,junkc)
zero_image=.true.
end function



!-------------------------------------------------!
!                                                 !
!            GRAB/STORE ROUTINES                  !
!                                                 !
!-------------------------------------------------!

subroutine store_current_image_tag(r)
character(len=*)  :: r
rtag=r
end subroutine

subroutine grab_current_image_tag(r)
character(len=*) :: r
r=rtag
end subroutine



subroutine store_current_image_space(r)
type(sep3d) :: r
rsep=r
end subroutine

subroutine grab_current_image_space(r)
type(sep3d) :: r
r=rsep
end subroutine

logical function distribute_image(tag_small,sep_small,tag_big,sep_big)
character(len=*) :: tag_small,tag_big
type(sep3d)      :: sep_small,sep_big
integer          :: i3,i4,i5,fwind(6),nwind(6),nloop,nbig,esize,i
real, allocatable :: bigr(:,:),bigr2(:,:)
complex, allocatable :: bigc(:,:),bigc2(:,:)
logical             :: big,comp,adj,add,copy
character(len=128)   :: temp1,temp2,operation
distribute_image=.false.
call grab_current_adj(adj)


call my_operation(operation)
call grab_current_add(add)

copy=.true.
if(operation(1:4)=="migr" .or. .not. add)  copy=.false.


if(adj .and. sep_num_thread()>1) then !This is a sectioned dataset
  call from_param("R",temp1)
  call parcat(temp1,2) 
  call from_param("Rtemp",temp2,temp1)
  if(my_restart()) then
     call init_sep3d(temp2,sep_big,"SCRATCH")
  else
    if(.not. image_distrib_tag(sep_small,sep_big,.true.)) then
       write(0,*) "trouble create distirbuted image tag"
       return
    end if
    call auxinout(temp1)
    call sep3d_write_description(temp1,sep_big)
    if(copy) then
      tag_big=temp1
      nwind=1;fwind=0; nwind(1:3)=sep_small%n(1:3); nbig=sep_small%n(2)
      if(product(nwind(1:3)) >2000000) then
        big=.false.; nloop=sep_small%n(3); nwind(3)=1;
      else
        big=.true.; nloop=1; nwind(3)=sep_small%n(3);nbig=nbig*sep_small%n(3)
      end if
      if(sep_small%data_format(1:1)=="F") then
        comp=.false. ; allocate(bigr(nwind(1),nbig),bigc(1,1));
        esize=4
      else
        comp=.true. ; allocate(bigc(nwind(1),nbig),bigr(1,1));
        esize=8
      end if
      fwind(6)=sep_thread_num()
      if(.not. init_print_progress(0,sep_big%n(5))) then
        write(0,*) "trouble initializing print progress"
        return 
      end if
      do i5=1,sep_big%n(5)
        call print_progress("Distributing image")
        fwind(5)=i5-1
        do i4=1,sep_big%n(4)
          fwind(4)=i4-1
          do i3=1,nloop
            fwind(3)=i3-1
            call sep3d_set_window(sep_big,nwind=nwind,fwind=fwind)
            call sep3d_set_window(sep_small,nwind=nwind,fwind=fwind)
            if(comp) then
              if(sep_thread_num()==0) then
                if(.not. sep3d_read_data(tag_small,sep_small,bigc,&
                  local_tag=.true.)) then
                   write(0,*) "trouble reading data",fwind
                   return
                 end if
              end if
              if(.not. sep3d_spread_data(sep_big,bigc)) then
                  write(0,*) "trouble combining data ", fwind
                  return
              end if
              if(.not. sep3d_write_data(tag_big,sep_big,bigc,fwind=fwind,&
                 nwind=nwind)) then
                  write(0,*) "trouble writing ", fwind
                  return
              end if
            else
              if(sep_thread_num()==0) then
                if(.not. sep3d_read_data(tag_small,sep_small,bigr,&
                  local_tag=.true.)) then
                   write(0,*) "trouble reading data",fwind
                   return
                 end if
              end if
              if(.not. sep3d_spread_data(sep_big,bigr)) then
                  write(0,*) "trouble combining data ", fwind
                  return
              end if
              if(.not. sep3d_write_data(tag_big,sep_big,bigr,fwind=fwind,&
                 nwind=nwind)) then
                  write(0,*) "trouble writing ", fwind
                  return
              end if
            end if
          end do
        end do
      end do
      deallocate(bigr,bigc)
    end if
  end if
else
  tag_big=tag_small
  sep_big=sep_small
end if
distribute_image=.true.
end function





logical function combine_image(big_tag,big_sep,small_tag,small_sep)
character(len=*) :: big_tag,small_tag
type(sep3d)      :: big_sep,small_sep
integer          :: i3,i4,i5,fwind(6),nwind(6),nloop,nbig
real, allocatable :: bigr(:,:),bigr2(:,:)
complex, allocatable :: bigc(:,:),bigc2(:,:)
logical             :: big,comp,adj
combine_image=.false.
call grab_current_adj(adj)

if(sep3d_ndims(big_sep) >5 .and. adj) then !This is a sectioned dataset
  if(.not. init_print_progress(0,big_sep%n(5))) then
    write(0,*) "trouble initializing print progress"
    return
  end if

  nwind=1;fwind=0; nwind(1:3)=big_sep%n(1:3); nbig=big_sep%n(2)
  if(product(nwind(1:3)) >2000000) then
    big=.false.; nloop=big_sep%n(3); nwind(3)=1;
  else
    big=.true.; nloop=1; nwind(3)=big_sep%n(3);nbig=nbig*big_sep%n(3)
  end if
  if(big_sep%data_format(1:1)=="F") then
    comp=.false. ; allocate(bigr(nwind(1),nbig),bigc(1,1));
    comp=.false. ; allocate(bigr2(nwind(1),nbig),bigc2(1,1));
  else
    comp=.true. ; allocate(bigc(nwind(1),nbig),bigr(1,1));
    comp=.true. ; allocate(bigc2(nwind(1),nbig),bigr2(1,1));
  end if
  fwind(6)=sep_thread_num()
  do i5=1,small_sep%n(5)
    fwind(5)=i5-1
    call print_progress("Combining image")
    do i4=1,small_sep%n(4)
      fwind(4)=i4-1
      do i3=1,nloop
        fwind(3)=i3-1
        call sep3d_set_window(small_sep,nwind=nwind,fwind=fwind)
        call sep3d_set_window(big_sep,nwind=nwind,fwind=fwind)
        if(comp) then
          if(.not. sep3d_read_data(big_tag,big_sep,bigc)) then
              write(0,*) "trouble reading ", fwind
              return
          end if
          if(.not. sep3d_compress_data(big_sep,small_sep,bigc,bigc2)) then
              write(0,*) "trouble combining data ", fwind
              return
          end if
          if(sep_thread_num()==0) then
            if(.not. sep3d_write_data(small_tag,small_sep,bigc2,local_tag=.true.,nwind=nwind,fwind=fwind)) then
               write(0,*) "trouble writing data",fwind
               return
            end if
          end if
        else
          if(.not. sep3d_read_data(big_tag,big_sep,bigr)) then
              write(0,*) "trouble reading ", fwind
              return
          end if
          if(.not. sep3d_compress_data(big_sep,small_sep,bigr,bigr2)) then
              write(0,*) "trouble combining data ", fwind
              return
          end if
          if(sep_thread_num()==0) then
            if(.not. sep3d_write_data(small_tag,small_sep,bigr2,local_tag=.true.,nwind=nwind,fwind=fwind)) then
               write(0,*) "trouble writing data",fwind
               return
            end if
          end if
        end if
      end do
    end do
  end do
  deallocate(bigr,bigr2,bigc,bigc2)
end if
combine_image=.true.
end function





end module
