program OWmig2d

use sep 
!use OW_io
!use OW_vref
use OW_type
use fft_data_mod

integer						:: ixs,iw,iz,ikx,iref,ixw,ish
integer						:: nref,ni,ihso,isp
real						:: xmin,xmax,ymin,ymax,zmin,zmax
real						:: sign, pi
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(:,:,:)

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

 call sep_init()
 call sep_begin_prog()

!
! 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%xm%n);	call from_aux("vel","d1",v%xm%d);	call from_aux("vel","o1",v%xm%o)
 call from_aux("vel","n3",v%z%n);	call from_aux("vel","d3",v%z%d);	call from_aux("vel","o3",v%z%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.)

 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=(xmax-xmin)/v%xm%d + 1
 img%xm%d=v%xm%d
 img%xm%o=xmin

 img%z%n=v%z%n
 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)  


 allocate(recv(rec%h%n,rec%w%n,rec%x%n))
 allocate(source(sou%h%n,sou%w%n,sou%x%n))
 allocate(velint(v%xm%n,v%z%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(v%z%n,rec%x%n,rec%h%n))

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;
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
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
fft_1d_idx = init_fft_1d(kxm%n)
!--------------------------------------------------
! 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       
    ni=floor((owrk-v%xm%o)/v%xm%d)+1;
    vel=0;
    vel(:,:)=velint(:,ni:ni+rec%h%n-1);
!
    do iz=1,v%z%n
        if (nref .eq. 1) then
           vref(nref,iz,ixs)=sum(vel(iz,:))/rec%h%n
        else
           minv=minval(vel(iz,:))
           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 do 
end do 
!--------------------------------------------------
!
!
write(0,*) "DOWNWARD CONTINUES SOURCES AND RECEIVERS"
write(0,*) "========================================"
!
! sws=zeros(nw,nz,nkx);swr=sws;
 allocate(ws(rec%h%n),wr(rec%h%n))
 allocate(gs(rec%h%n,nref),gr(rec%h%n,nref))
 allocate(gs0(rec%h%n),gr0(rec%h%n),sqa0(rec%h%n))
!--------------------------------------------------
! shot position in the acq.geom.
!isp = ixpad - floor(rec%h%o/rec%h%d) + 1
isp = abs(floor(rec%h%o/rec%h%d)) + 1
 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)
!--------------------------------------------------
! Migrating
!
shot:do ixs=1,rec%x%n

 call sreed_window("in",3,(/rec%h%n,rec%w%n,rec%x%n/),&
                  (/rec%h%n,rec%w%n,rec%x%n/),(/0,0,ixs-1/),&
                  (/1,1,1/),8,recv)

    GG=0.
! workspace origin 
    owrk=(ixs-1)*rec%x%d+rec%x%o+rec%h%o
! grab velocity in the workspace       
    ni=floor((owrk-v%xm%o)/v%xm%d)+1;
    vel=0;
    vel(:,:)=velint(:,ni:ni+rec%h%n-1);
!      
    ws=0;wr=0
    gs=0;gr=0
    gs0=0;gr0=0
!!! !$OMP parallel do private(inode,iw,ixg,ixp,iz,iv,ik,nref,sqa,tmp_wfld,G1,vtmp,green,green1,vref,ind_vref,iixs,iixr,xxs,xxr)   

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

        inode=omp_get_thread_num();

        ws(isp)=source(1,iw,1) ! to be modified to handle generalized sources
        wr=recv(:,iw,ixs)
        depth:do iz=2,img%z%n

            call fft_1d_data(fft_1d_idx,  -1,  kxm%n,ws)
            call fft_1d_data(fft_1d_idx,  -1,  kxm%n,wr)

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

            vxz:if (minv .eq. maxv) then
                gs0=ws;  gr0=wr;
!      Operations in the wavenumber domain
!                wavenumber:do ik = 1:kxm%n
!      phase shift
!                   sqa0 = rec%w%all(iw)**2 / minv**2 - kxm%all(ik)**2
!                   evanescent:if (sqa0 > 0) then
!                       gs0(ik)=gs0(ik)*exp(complex(0,-sign*dz*sqrt(sqa0)))
!                       gr0(ik)=gr0(ik)*exp(complex(0,sign*dz*sqrt(sqa0)))
!                   else
!                       gs0(ik)=gs0(ik)*exp(-sqrt(-sqa0)*dz)
!                       gr0(ik)=gr0(ik)*exp(-sqrt(-sqa0)*dz)
!                   end if evanescent
!                end do wavenumber

                   sqa0 = rec%w%all(iw)**2 / minv**2 - kxm%all**2
                   evanescent1:where (sqa0 > 0)
                       gs0=gs0*exp(cmplx(0,-sign*dz*sqrt(sqa0)))
                       gr0=gr0*exp(cmplx(0,sign*dz*sqrt(sqa0)))
                   elsewhere
                       gs0=gs0*exp(-sqrt(-sqa0)*dz)
                       gr0=gr0*exp(-sqrt(-sqa0)*dz)
                   end where evanescent1

! Operations in x            

                   call fft_1d_data(fft_1d_idx,  1,  kxm%n,gs0)
                   call fft_1d_data(fft_1d_idx,  1,  kxm%n,gr0)

!      split-step
                   gs0 = gs0 * exp(cmplx(0,-sign*dz*rec%w%all(iw))*(1./vel(iz,:)-1./minv))
                   gr0 = gr0 * exp(cmplx(0, sign*dz*rec%w%all(iw))*(1./vel(iz,:)-1./minv))
                   ws=gs0;   wr=gr0
            else;                                                                                       ! v(z) condition
                velref: do iref=1,nref                                                                        ! reference velocities                 
                    gs(:,iref)=ws;  gr(:,iref)=wr; 
!                    do ikx = 1,nkx;                                                                         ! wavenumber loop
!!      phase shift
!                        sqa0 = rec%w%all(iw)**2 / vref(iref,iz,ixs)**2 - kxall(ikx)**2;
!                        if (sqa0 > 0) then
!                            gs(ikx,iref)=gs(ikx,iref)*exp(cmplx(0,-sign*dz*sqrt(sqa0)));
!                            gr(ikx,iref)=gr(ikx,iref)*exp(complex(0,sign*dz*sqrt(sqa0)));
!                        else
!                            gs(ikx,iref)=0.;
!                            gr(ikx,iref)=0.;
!                        end if
!                    end do                                                                                    ! end wavenumber loop

                      sqa0 = rec%w%all(iw)**2 / vref(iref,iz,ixs)**2 - kxall**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            

                      call fft_1d_data(fft_1d_idx,  1,  kxm%n,gs(:,iref))
                      call fft_1d_data(fft_1d_idx,  1,  kxm%n,gr(:,iref))

!      split step
                      gs(:,iref) = gs(:,iref) * exp(cmplx(0,-sign*dz*rec%w%all(iw))*(1./vel(iz,:)-1./vref(iref,iz,ixs)));
                      gr(:,iref) = gr(:,iref) * exp(cmplx(0, sign*dz*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:nkx;
!                    do iref=1:nref-1;
!                        if vref(iref,iz,ixs) ~= vref(iref+1,iz,ixs);
!                            if vel(iz,ix) >= vref(iref,iz,ixs) .and. vel(iz,ix) < vref(iref+1,iz,ixs);
!                                f=(vel(iz,ix) - vref(iref,iz,ixs))/vel(iz,ix); g=1-f;
!                                ws(iz,ix)=g*gs(ix,iref)+f*gs(ix,iref+1);
!                                wr(iz,ix)=g*gr(ix,iref)+f*gr(ix,iref+1);
!!                                ws(iz,ix)=gs(ix,iref);
!!                                wr(iz,ix)=gr(ix,iref);
!                            elseif vel(iz,ix) >= vref(nref,iz,ixs);
!                                ws(iz,ix)=gs(ix,nref);
!                                wr(iz,ix)=gr(ix,nref);
!                            end if
!                        else;
!                            ws(iz,ix)=gs(ix,iref);
!                            wr(iz,ix)=gr(ix,iref);
!                        end if                            
!                    end do
!                end do
           end if vxz
!
! Imaging Condition
!
            do ish=1,rec%x%n
                ihso=floor(((ish-1)*rec%x%d+rec%x%o)/(v%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
                        cycle
                    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
        write(0,*)  "DONE FREQUENCY # ",iw," OUT OF ",rec%w%n," / SHOT # ",rec%x%n
!        sws(iw,:,:)=ws
!        swr(iw,:,:)=wr
    end do omega

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

    write(0,*) " DONE SHOT # ",ixs," OUT OF ",rec%x%n," in "," ?? "," min"
end do shot                                                                                                    ! end shot loop

call srite("out",real(image),4*img%z%n*img%sh%n*img%xm%n)
!sws(nw+1:nt,:,:)=0;
!swr(nw+1:nt,:,:)=0;
!sws=real(ifft(conj(sws)));
!swr=real(ifft((swr)));
!a=zeros(floor(nt/5)+1,nz,nkx);b=a;
!do i=1:5:nt;
!    b((i-1)/5+1,:,:)=swr(i,:,:);
!    a((i-1)/5+1,:,:)=sws(i,:,:);
!end do


 call sep_end_prog()
 end program OWmig2d
