module OW_pspi3

use OW_type
use OW_parms

implicit none

contains

!--------------------------------------------------------
subroutine owpspi3(iw,isp,planF,planI,recv,source,vref,vel,GG)

integer						:: ixs,iz,ikx,iref,ixw,ish,ihso
integer						:: iw,isp
real						:: sqa,minv, maxv
real						:: vref(:,:),vel(:,:)
real, allocatable				:: sqa0(:)
complex, allocatable				:: ws(:),wr(:),gs0(:),gr0(:),gs(:,:),gr(:,:)
complex 					:: recv(:,:),source(:,:,:),GG(:,:,:)
complex, allocatable				:: tfield(:)

integer*8					:: planF, planI

!-------------------------------------------------------------------

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

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

        depth:do iz=2,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, v%z%d*sqrt(sqa)))
!                       gr0(ik)=gr0(ik)*exp(cmplx(0,-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

                   sqa0 = rec%w%all(iw)**2 / minv**2 - kxm%all**2
                   evanescent1:where (sqa0 > 0)
                       gs0=gs0*exp(cmplx(0, v%z%d*sqrt(sqa0)))
                       gr0=gr0*exp(cmplx(0,-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
                velref: do iref=1,nref

                    gs(:,iref)=ws;  gr(:,iref)=wr; 
!                    do ikx = 1,kxm%n
!      phase shift
!                        sqa = rec%w%all(iw)**2 / vref(iref,iz)**2 - kxm%all(ikx)**2;
!                        if (sqa > 0) then
!                            gs(ikx,iref)=gs(ikx,iref)*exp(cmplx(0, v%z%d*sqrt(sqa)));
!                            gr(ikx,iref)=gr(ikx,iref)*exp(cmplx(0,-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)**2 - kxm%all**2;
                      evanescent2:where (sqa0 > 0)
                          gs(:,iref)=gs(:,iref)*exp(cmplx(0, v%z%d*sqrt(sqa0)))
                          gr(:,iref)=gr(:,iref)*exp(cmplx(0,-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, v%z%d*rec%w%all(iw))*(1./vel(iz,:)-1./vref(iref,iz)));
                      gr(:,iref) = gr(:,iref) * exp(cmplx(0,-v%z%d*rec%w%all(iw))*(1./vel(iz,:)-1./vref(iref,iz)));
                end do velref

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

!                do ix=1,kxm%n
!                    do iref=1,nref-1
!                        if (vref(iref,iz) .ne. vref(iref+1,iz)) then
!                            if (vel(iz,ix) >= vref(iref,iz) .and. vel(iz,ix) < vref(iref+1,iz)) then
!                                f=(vel(iz,ix) - vref(iref,iz))/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)) 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=floor(((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
deallocate(ws)
deallocate(wr)
deallocate(gs)
deallocate(gr)
deallocate(gs0)
deallocate(gr0)
deallocate(sqa0)
deallocate(tfield)

end subroutine
end module
