program OWmig2d1

use sep 
use OW_type

integer						:: time0(8), time1(8)
real						:: start_time, finish_time

integer						:: ixs,iw,iz,ikx,iref,ixw,ish,ino,k
integer						:: nref,ni,nv,ihso,isp,inode,nnode
real						:: xmin,xmax,ymin,ymax,zmin,zmax,vsal
real						:: sign, pi,sqa
real						:: minv, maxv, dvel, owrk
real, allocatable				:: vref(:,:,:),vel(:,:),velint(:,:),sqa0(:)
complex, allocatable				:: ws(:),wr(:),gs0(:),gr0(:),gs(:,:),gr(:,:)
complex, allocatable				:: recv(:,:),source(:,:,:),image(:,:,:),GG(:,:,:)!,:)
complex, pointer				:: tfield(:)

type (axis) 					:: kxm, h, w, x, xm, z, sh
type (dat_grid)					:: rec, sou
type (mod_grid)					:: img
type (velo)					:: v
!type (fft_type)     				:: fft_1d_idx

integer, external 				:: omp_get_thread_num
integer, external 				:: omp_get_num_threads

integer, parameter   				:: FFTW_R2HC=0,FTW_HC2R=1,FFTW_DHT=2,FFTW_REDFT00=3,FFTW_REDFT01=4
integer, parameter   				:: FFTW_REDFT10=5,FFTW_REDFT11=6,FFTW_RODFT00=7,FFTW_RODFT01=8,FFTW_RODFT10=9
integer, parameter   				:: FFTW_RODFT11=10,FFTW_FORWARD=-1,FFTW_BACKWARD=+1,FFTW_DESTROY_INPUT=1
integer, parameter   				:: FFTW_UNALIGNED=2,FFTW_CONSERVE_MEMORY=4,FFTW_EXHAUSTIVE=8,FTW_PRESERVE_INPUT=16
integer, parameter   				:: FFTW_PATIENT=32,FFTW_ESTIMATE=64,FFTW_ESTIMATE_PATIENT=128
integer, parameter   				:: FFTW_BELIEVE_PCOST=256,FFTW_DFT_R2HC_ICKY=512,FFTW_NONTHREADED_ICKY=1024
integer, parameter   				:: FFTW_NO_BUFFERING=2048,FFTW_NO_INDIRECT_OP=4096,FFTW_ALLOW_LARGE_GENERIC=8192
integer, parameter   				:: FFTW_NO_RANK_SPLITS=16384,FFTW_NO_VRANK_SPLITS=32768,FFTW_NO_VRECURSE=65536
 
!integer*8, pointer,private   			:: planF,planI
!integer*8, private	   			:: planF,planI
integer*8		   			:: planF,planI  

 call sep_init()
 call sep_begin_prog()

start_time=0.;finish_time=0.
call date_and_time(values=time0)
start_time = time0 (5) * 3600 + time0 (6) * 60 + time0 (7) + 0.001 * time0 (8)

!
! IO
!

 call from_history("n1",rec%h%n);	call from_history("d1",rec%h%d);	call from_history("o1",rec%h%o)
 call from_history("n3",rec%w%n);	call from_history("d3",rec%w%d);	call from_history("o3",rec%w%o)
 call from_history("n4",rec%x%n);	call from_history("d4",rec%x%d);	call from_history("o4",rec%x%o)  

 call from_aux("sou","n1",sou%h%n);	call from_aux("sou","d1",sou%h%d);	call from_aux("sou","o1",sou%h%o)
 call from_aux("sou","n3",sou%w%n);	call from_aux("sou","d3",sou%w%d);	call from_aux("sou","o3",sou%w%o)
 call from_aux("sou","n4",sou%x%n);	call from_aux("sou","d4",sou%x%d);	call from_aux("sou","o4",sou%x%o)  

 call from_aux("vel","n1",v%z%n);	call from_aux("vel","d1",v%z%d);	call from_aux("vel","o1",v%z%o)
 call from_aux("vel","n2",v%xm%n);	call from_aux("vel","d2",v%xm%d);	call from_aux("vel","o2",v%xm%o)

 if (sou%w%n .ne. rec%w%n) then
    call seperr("# source freq. .ne. # receiver freq.")
 end if
 if (sou%w%d .ne. rec%w%d) then
    call seperr("source freq.interval .ne. receiver freq.interval")
 end if
 if (sou%w%o .ne. rec%w%o) then
    call seperr("source freq.origin .ne. receiver freq.origin")
 end if

 !!  get parameters
  call from_param("nrefvel",nref) 
  call from_param("xmin",xmin) 
  call from_param("xmax",xmax)
  call from_param("ymin",ymin,0.)
  call from_param("ymax",ymax,0.)
  call from_param("zmin",zmin,0.)
  call from_param("zmax",zmax)
  call from_param("nh",img%sh%n,1)
  call from_param("oh",img%sh%o,0.)
  call from_param("dh",img%sh%d,0.)
  call from_param("vsal",vsal)

 if (xmin < v%xm%o ) call seperr("Min.x-image space < Min.x-velocity")
 if (xmax > v%xm%o+(v%xm%n-1)*v%xm%d ) call seperr("Max.x-image space > Max.x-velocity")

 if (zmin < v%xm%o ) call seperr("Min.z-image space < Min.z-velocity")
 if (zmax > v%z%o+(v%z%n-1)*v%z%d ) call seperr("Max.z-image space > Max.z-velocity")

 if (rec%h%o + rec%x%o < xmin ) call seperr("Min.x-receiver < Min.x-image space")
 if (rec%h%o+(rec%h%n-1)*rec%h%d + rec%x%o+(rec%x%n-1)*rec%x%d > xmax ) call seperr("Max.x-receiver < Max.x-image space")

 if (sou%h%o + sou%x%o < xmin ) call seperr("Min.x-source < Min.x-image space")
 if (sou%h%o+(sou%h%n-1)*sou%h%d + sou%x%o+(sou%x%n-1)*sou%x%d > xmax ) call seperr("Max.x-source < Max.x-image space")

 if (rec%h%d .ne. v%xm%d) call seperr("x-offset interval .ne. velocity x-interval")

 img%xm%n=floor((xmax-xmin)/v%xm%d)+1.5
 img%xm%d=v%xm%d
 img%xm%o=xmin

 img%z%n=floor((zmax-zmin)/v%z%d)+1.5
 img%z%o=zmin
 img%z%d=v%z%d

 call to_history("n1",img%z%n) ; call to_history("d1",img%z%d) ; call to_history("o1",img%z%o)
 call to_history("n2",img%sh%n); call to_history("d2",img%sh%d); call to_history("o2",img%sh%o)
 call to_history("n3",img%xm%n); call to_history("d3",img%xm%d); call to_history("o3",img%xm%o)  
 call to_history("n4",1); call to_history("d4",1); call to_history("o4",1)
 call to_history("n5",1); call to_history("d5",1); call to_history("o5",1)  
 call to_history("esize",4)

nnode=omp_get_num_threads()

 allocate(recv(rec%h%n,rec%w%n))
 allocate(source(sou%h%n,sou%w%n,sou%x%n))
 allocate(velint(v%z%n,v%xm%n))
 allocate(image(img%z%n,img%sh%n,img%xm%n))
 allocate(vref(nref,v%z%n,rec%x%n))
 allocate(vel(v%z%n,rec%h%n))
 allocate(GG(img%z%n,img%sh%n,rec%h%n))!,nnode))

sign=-1;
pi=4*atan(1.)

!--------------------------------------------------
! Frequency vector
!
! no zero frequency component
rec%w%o   = 2*pi * rec%w%o
rec%w%d   = 2*pi * rec%w%d
allocate(rec%w%all(rec%w%n))
do iw = 1,rec%w%n
   rec%w%all(iw) = rec%w%o + (iw - 1) * rec%w%d
end do
!--------------------------------------------------
! Wavenumber vector
!
kxm%n = rec%h%n
kxm%d = 1./(kxm%n*rec%h%d)
kxm%o = 1./(2*rec%h%d)-kxm%d
allocate(kxm%all(kxm%n))
do ikx=1,kxm%n
   kxm%all(ikx) = mod(kxm%o+(ikx-1)*kxm%d,1./rec%h%d)-kxm%o
end do
kxm%all = 2*pi * kxm%all

!--------------------------------------------------
! Setting FFTW plans
!

allocate(tfield(kxm%n))
call sfftw_plan_dft(planF,1,kxm%n,tfield,tfield,fftW_FORWARD,fftW_PATIENT)
call sfftw_plan_dft(planI,1,kxm%n,tfield,tfield,fftw_BACKWARD,fftW_PATIENT)

!--------------------------------------------------
! Reading source fctn and velocity

 call sreed("sou",source,sou%h%n*sou%w%n*sou%x%n*8)
 call sreed("vel",velint,v%z%n*v%xm%n*4)

!--------------------------------------------------
! Pre-compute vref
!

vref=0
do ixs=1,rec%x%n
    owrk=(ixs-1)*rec%x%d+rec%x%o+rec%h%o
    if (owrk < img%xm%o .or. owrk > img%xm%o+img%xm%d*(img%xm%n-1)) call seperr('WORKING SPACE GREATER THAN IMAGE SPACE FOR SHOT # ',ixs)
! grab velocity in the workspace       
    nv=floor((owrk-v%xm%o)/v%xm%d)+1.5
    vel=0;
    vel=velint(:,nv:nv+rec%h%n-1);
!
    do iz=1,v%z%n
!        k=count(vel(iz,:) < vmax))
!        where(vel(iz,:) < vmax)
        if (nref .eq. 1) then
           vref(nref,iz,ixs)=sum(vel(iz,:))/rec%h%n
        else
           minv=minval(vel(iz,:))
           if (maxval(vel(iz,:)) >= vsal) then
               maxv=maxval(vel(iz,:),1,vel(iz,:) < vsal)
               dvel=(maxv-minv)/(nref-2) 
               do iref=1,nref-1
                  vref(iref,iz,ixs)=minv+(iref-1)*dvel
               end do 
               vref(nref,iz,ixs)=vsal
           else
               maxv=maxval(vel(iz,:))
               dvel=(maxv-minv)/(nref-1) 
               do iref=1,nref
                  vref(iref,iz,ixs)=minv+(iref-1)*dvel
               end do 
           end if 
        end if
    end do 
end do 
if(exist_file("vref")) then
 call to_history("n1",nref,"vref"); call to_history("d1",1.,"vref"); call to_history("o1",0.,"vref")
 call to_history("n2",img%z%n,"vref") ; call to_history("d2",img%z%d,"vref") ; call to_history("o2",img%z%o,"vref")
 call to_history("n3",rec%x%n,"vref"); call to_history("d3",rec%x%d,"vref"); call to_history("o3",rec%x%o,"vref")  
 call to_history("esize",4)
 call srite("vref",vref,4*nref*img%z%n*rec%x%n)
end if 

call date_and_time(values=time1)
finish_time = time1 (5) * 3600 + time1 (6) * 60 + time1 (7) + 0.001 * time1 (8)
write(0,'(8x,a,f10.2,a)')' Init.phase - Total elapsed time:', (finish_time - start_time)/60./60.,' (hrs)'


!--------------------------------------------------
!
!
write(0,*) "DOWNWARD CONTINUES SOURCES AND RECEIVERS"
write(0,*) "========================================"
!
!--------------------------------------------------
! shot position in the acq.geom.
isp = abs(floor(rec%h%o/rec%h%d)) + 1.5

!--------------------------------------------------
! Migrating
!

shot:do ixs=1,rec%x%n

start_time=0.;finish_time=0.
call date_and_time(values=time0)
start_time = time0 (5) * 3600 + time0 (6) * 60 + time0 (7) + 0.001 * time0 (8)

  call sreed("in",recv,8*rec%h%n*rec%w%n)

    GG=0.
! workspace origin 
    owrk=(ixs-1)*rec%x%d+rec%x%o+rec%h%o
! grab velocity in the workspace       
    nv = floor((owrk-v%xm%o)/v%xm%d)+1.5
    ni = nv-floor((img%xm%o-v%xm%o)/img%xm%d)+1.5

    write(0,*) owrk,nv,ni

    vel=0;
    vel(:,:)=velint(:,nv:nv+rec%h%n-1);

    omega:do iw = 1,rec%w%n

        if (allocated(ws)) deallocate(ws);	allocate(ws(rec%h%n));		ws=0
        if (allocated(wr)) deallocate(wr);	allocate(wr(rec%h%n));		wr=0
!        if (allocated(tfield)) deallocate(tfield);	allocate(tfield(rec%h%n));	
        if (allocated(gs)) deallocate(gs);	allocate(gs(rec%h%n,nref));	gs=0
        if (allocated(gr)) deallocate(gr);	allocate(gr(rec%h%n,nref));	gr=0
        if (allocated(gs0)) deallocate(gs0);	allocate(gs0(rec%h%n));		gs0=0
        if (allocated(gr0)) deallocate(gr0);	allocate(gr0(rec%h%n));		gr0=0
        if (allocated(sqa0)) deallocate(sqa0);	allocate(sqa0(rec%h%n));	sqa0=0
!        if (allocated(vref)) deallocate(vref);	allocate(vref(nref));		vref=0

        ws(isp)=source(1,iw,1) ! to be modified to handle generalized sources
        wr=recv(:,iw)

        depth:do iz=1,img%z%n

            tfield=ws; call sfftw_execute(planF); ws=tfield/sqrt(1.*kxm%n)
            tfield=wr; call sfftw_execute(planF); wr=tfield/sqrt(1.*kxm%n)

            minv=minval(vel(iz,:)); maxv=maxval(vel(iz,:))

            vxz:if (minv .eq. maxv) then
                gs0=ws;  gr0=wr;
!      Operations in the wavenumber domain
!                wavenumber1:do ik = 1,kxm%n
!      phase shift
!                   sqa = rec%w%all(iw)**2 / minv**2 - kxm%all(ik)**2
!                   evanescent1:if (sqa > 0) then
!                       gs0(ik)=gs0(ik)*exp(cmplx(0,-sign*v%z%d*sqrt(sqa)))
!                       gr0(ik)=gr0(ik)*exp(cmplx(0,sign*v%z%d*sqrt(sqa)))
!                   else
!                       gs0(ik)=gs0(ik)*exp(-sqrt(-sqa)*v%z%d)
!                       gr0(ik)=gr0(ik)*exp(-sqrt(-sqa)*v%z%d)
!                   end if evanescent1
!                end do wavenumber1
        if (mod(iw,10)==1) write(2,*)  "2 ",iw,iz,maxval(cabs(gs0)),maxval(cabs(gr0))
                   sqa0 = rec%w%all(iw)**2 / minv**2 - kxm%all**2
                   evanescent1:where (sqa0 > 0)
                       gs0=gs0*exp(cmplx(0,-sign*v%z%d*sqrt(sqa0)))
                       gr0=gr0*exp(cmplx(0,sign*v%z%d*sqrt(sqa0)))
                   elsewhere
                       gs0=gs0*exp(-sqrt(-sqa0)*v%z%d)
                       gr0=gr0*exp(-sqrt(-sqa0)*v%z%d)
                   end where evanescent1

! Operations in x            

                   tfield=gs0; call sfftw_execute(planI); ws=tfield/sqrt(1.*kxm%n)
                   tfield=gr0; call sfftw_execute(planI); wr=tfield/sqrt(1.*kxm%n)

            else;                                                                                       ! v(z) condition
                velref: do iref=1,nref

                    gs(:,iref)=ws;  gr(:,iref)=wr; 
!                    do ikx = 1,kxm%n;                                                                         ! wavenumber loop
!      phase shift
!                        sqa = rec%w%all(iw)**2 / vref(iref,iz,ixs)**2 - kxm%all(ikx)**2;
!                        if (sqa > 0) then
!                            gs(ikx,iref)=gs(ikx,iref)*exp(cmplx(0,-sign*v%z%d*sqrt(sqa)));
!                            gr(ikx,iref)=gr(ikx,iref)*exp(cmplx(0,sign*v%z%d*sqrt(sqa)));
!                        else
!                            gs(ikx,iref)=gs(ikx,iref)*exp(-sqrt(-sqa)*v%z%d)
!                            gr(ikx,iref)=gr(ikx,iref)*exp(-sqrt(-sqa)*v%z%d)
!                        end if
!                    end do                                                                                    ! end wavenumber loop

                      sqa0 = rec%w%all(iw)**2 / vref(iref,iz,ixs)**2 - kxm%all**2;
                      evanescent2:where (sqa0 > 0)
                          gs(:,iref)=gs(:,iref)*exp(cmplx(0,-sign*v%z%d*sqrt(sqa0)))
                          gr(:,iref)=gr(:,iref)*exp(cmplx(0, sign*v%z%d*sqrt(sqa0)))
                      elsewhere
                          gs(:,iref)=gs(:,iref)*exp(-sqrt(-sqa0)*v%z%d)
                          gr(:,iref)=gr(:,iref)*exp(-sqrt(-sqa0)*v%z%d)
                      end where evanescent2

! Operations in x            

                      tfield=gs(:,iref); call sfftw_execute(planI); gs(:,iref)=tfield/sqrt(1.*kxm%n)
                      tfield=gr(:,iref); call sfftw_execute(planI); gr(:,iref)=tfield/sqrt(1.*kxm%n)


!      split step
                      gs(:,iref) = gs(:,iref) * exp(cmplx(0,-sign*v%z%d*rec%w%all(iw))*(1./vel(iz,:)-1./vref(iref,iz,ixs)));
                      gr(:,iref) = gr(:,iref) * exp(cmplx(0, sign*v%z%d*rec%w%all(iw))*(1./vel(iz,:)-1./vref(iref,iz,ixs)));
                end do velref

                do iref=1,nref-1
                interpwvfld: where (vel(iz,:) >= vref(iref,iz,ixs) .and. vel(iz,:) < vref(iref+1,iz,ixs))
                     ws=gs(:,iref)
                     wr=gr(:,iref)
                     end where interpwvfld
                end do
                where (vel(iz,:) >= vref(nref,iz,ixs))
                      ws=gs(:,nref)
                      wr=gr(:,nref)
                end where 

!                do ix=1,kxm%n
!                    do iref=1,nref-1
!                        if (vref(iref,iz,ixs) .ne. vref(iref+1,iz,ixs)) then
!                            if (vel(iz,ix) >= vref(iref,iz,ixs) .and. vel(iz,ix) < vref(iref+1,iz,ixs)) then
!                                f=(vel(iz,ix) - vref(iref,iz,ixs))/vel(iz,ix)
!                                g=1-f
!                                ws(ix)=g*gs(ix,iref)+f*gs(ix,iref+1)
!                                wr(ix)=g*gr(ix,iref)+f*gr(ix,iref+1)
!!                                ws(ix)=gs(ix,iref)
!!                                wr(ix)=gr(ix,iref)
!                            elseif (vel(iz,ix) >= vref(nref,iz,ixs)) then
!                                ws(ix)=gs(ix,nref)
!                                wr(ix)=gr(ix,nref)
!                            end if
!                        else
!                            ws(ix)=gs(ix,iref)
!                            wr(ix)=gr(ix,iref)
!                        end if                            
!                    end do
!                end do
           end if vxz

!
! Imaging Condition
!
            do ish=1,img%sh%n
                ihso=nint(((ish-1)*img%sh%d+img%sh%o)/(img%xm%d))
                do ixw=1,kxm%n
                    if (ixw-ihso < 1 .or. ixw-ihso > kxm%n .or. ixw+ihso < 1 .or. ixw+ihso > kxm%n) then
                    else
                         GG(iz,ish,ixw)=GG(iz,ish,ixw)+conjg(ws(ixw-ihso))*(wr(ixw+ihso))             ! correlation to image
                    end if
                end do
            end do
        end do depth
        if (mod(iw,10)==1) write(0,4)  "  DONE FREQUENCY # ",iw," OUT OF ",rec%w%n," / SHOT # ",ixs
4   format(a19,i3,a8,i3,a10,i3)
    end do omega

    image(:,:,ni:ni+kxm%n-1)=image(:,:,ni:ni+kxm%n-1)+GG;

    call date_and_time(values=time1)
    finish_time = time1 (5) * 3600 + time1 (6) * 60 + time1 (7) + 0.001 * time1 (8)
    write(0,5) " DONE SHOT # ",ixs," OUT OF ",rec%x%n," in ", (finish_time - start_time)/60," (min)"
5   format(a13,i3,a8,i3,a4,f6.2,a6)
end do shot

call srite("out",real(image),4*img%z%n*img%sh%n*img%xm%n)

 call sep_end_prog()
 end program OWmig2d1
