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


contains

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

init_image=.false.


if(.not. erase .and. 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 image struct"
    return
  end if
  created=.true.
end if
call store_current_image_space(myrsep)
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_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)

init_image_struct=.true.
end function


logical function image_tag(wsep,myrsep,spread)
type(sep3d) :: wsep,myrsep
logical :: myre,spread
image_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(5)="DEPTH"


call from_param("ray_par",ray_par,.false.)
if(ray_par) then
  myrsep%label(3)="P_X"
  myrsep%label(4)="P_Y"
  if(myrsep%n(3)>1) then
    call from_param("apx_o",myrsep%o(3),0.)
    call from_param("apx_d",myrsep%d(3),.3)
  end if
  if(myrsep%n(4)>1) then
    call from_param("apy_o",myrsep%o(4),0.)
    call from_param("apy_d",myrsep%d(4),.3)
  end if
else
  myrsep%label(3)="H_X"
  myrsep%label(4)="H_Y"
end if

call sep3d_set_sep3d(myrsep)

image_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(5),fw(5),ndim,nalloc,nloop,nh,iz,ielem
logical :: ddd,big
complex, allocatable :: junkc(:,:)
real, allocatable :: junkr(:,:)
integer :: npadx,npady
write_image_iz=.false.
 

big=.false.
call sep_open_lock(tag)
npadx= grab_current_npadx()
npady= grab_current_npady()


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

if(myrsep%n(1)*myrsep%n(2)*myrsep%n(3)< 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
  allocate(junkc(myrsep%n(1),nalloc),junkr(1,1))
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(.not. sep3d_read_data(tag,myrsep,junkc)) then
            write(0,*) "trouble reading image",iz
            return
          end if
          if(big) then
            junkc=reshape(wfld(npadx+1:npadx+nw(1),npady+1:npady+nw(2),:,i4,ielem),(/myrsep%n(1),nalloc/))+junkc
          else 
            junkc=wfld(npadx+1:npadx+nw(1),npady+1:npady+nw(2),i3,i4,ielem)+junkc
          end if
          ddd=sep3d_write_data(tag,myrsep,junkc,nwind=nw,fwind=fw)
        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(npadx+1:npadx+nw(1),npady+1:npady+nw(2),:,i4,ielem)),(/myrsep%n(1),nalloc/))+junkr
          else
            junkr=real(wfld(npadx+1:npadx+nw(1),npady+1:npady+nw(2),i3,i4,ielem))+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)
call sep_close_lock(tag)
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(5),fw(5),ndim,nalloc,nloop,nh,iz,ielem
logical :: ddd,big
integer :: npadx,npady
read_image_iz=.false.
 
npadx= grab_current_npadx()
npady= grab_current_npady()
call sep_open_lock(tag)

big=.false.


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

if(myrsep%n(1)*myrsep%n(2)*myrsep%n(3)< 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
          ddd=sep3d_read_data(tag,myrsep,junkc)
          if(big) then
            wfld(npadx+1:npadx+nw(1),npady+1:npady+nw(2),:,i4,ielem)=reshape(junkc,(/myrsep%n(1),myrsep%n(2),myrsep%n(3)/))
          else 
            wfld(npadx+1:npadx+nw(1),npady+1:npady+nw(2),i3,i4,ielem)=junkc
          end if
        else 
          ddd=sep3d_read_data(tag,myrsep,junkr)
          if(big) then
            wfld(npadx+1:npadx+nw(1),npady+1:npady+nw(2),:,i4,ielem)=&
                cmplx(reshape(junkr,(/myrsep%n(1),myrsep%n(2),myrsep%n(3)/)),0.)
          else
            wfld(npadx+1:npadx+nw(1),npady+1:npady+nw(2),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)
          return
        end if
      end do
    end do
call del_mem(size(junkr)+size(junkc))
deallocate(junkr,junkc)
call sep_close_lock(tag)
read_image_iz=.true.
end function


logical function zero_image(tag,myrsep)
type(sep3d) :: myrsep
character(len=*) :: tag
character(len=128) :: bb
real, allocatable :: junkr(:,:)
complex, allocatable :: junkc(:,:)
integer ::i3,i4,i5,nw(5),fw(5),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;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)
!          write(0,*) "ZXERO ING AMGE",trim(tag) 
        else 
          ddd=sep3d_write_data(tag,myrsep,junkr,nwind=nw,fwind=fw)
        end if
        if(.not. ddd) then
          deallocate(junkr,junkc)
          return
        end if
      end do
    end do
end do
call del_mem(size(junkc)+size(junkr))
deallocate(junkr,junkc)
!call from_param("R",bb)
!call seperr("")
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

end module
