module OW_pspi

use OWvel

implicit none

contains
!---------------------------------------------------------
subroutine owpspi(recv,source,vel,ws,wr)

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

sign=-1;
pi=4*atan(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(tfld)) deallocate(tfld);	allocate(tfld(rec%h%n));	wr=0
        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,img%z%n));	vref=0

        ws=0;wr=0
        gs=0;gr=0
        gs0=0;gr0=0

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

        depth:do iz=2,img%z%n
!--------------------------------------------------
! Compute vref
!
            call owvref(nref,vel(iz,:),vref,minv,maxv)

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

            vxz:if (minv .eq. maxv) then
                gs0=ws;  gr0=wr;
!      Operations in the wavenumber domain
                wavenumber1:do ikx = 1,kxm%n
!      phase shift
                   sqa = rec%w%all(iw)**2 / minv**2 - kxm%all(ikx)**2
                   evanescent1:if (sqa > 0) then
                       gs0(ikx)=gs0(ikx)*exp(cmplx(0,-sign*v%z%d*sqrt(sqa)))
                       gr0(ikx)=gr0(ikx)*exp(cmplx(0,sign*v%z%d*sqrt(sqa)))
                   else
                       gs0(ikx)=gs0(ikx)*exp(-sqrt(-sqa)*v%z%d)
                       gr0(ikx)=gr0(ikx)*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,-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            

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

            else;                                                                                       ! v(z) condition
                velref: do iref=1,nref                                                                        ! reference velocities                 
                    gs(:,iref)=ws;  gr(:,iref)=wr; 
                    do ikx = 1,kxm%n;                                                                         ! wavenumber loop
!      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,-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)**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            

                      tfld=gs(:,iref); call sfftw_execute(planI); gs(:,iref)=tfld/sqrt(1.*kxm%n)
                      tfld=gr(:,iref); call sfftw_execute(planI); gr(:,iref)=tfld/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)));
                      gr(:,iref) = gr(:,iref) * exp(cmplx(0, sign*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 ikx=1,kxm%n
                    do iref=1,nref-1
                        if (vref(iref,iz) .ne. vref(iref+1,iz)) then
                            if (vel(iz,ikx) >= vref(iref,iz) .and. vel(iz,ikx) < vref(iref+1,iz)) then
                                f=(vel(iz,ikx) - vref(iref,iz))/vel(iz,ikx)
                                g=1-f
                                ws(ikx)=g*gs(ikx,iref)+f*gs(ikx,iref+1)
                                wr(ikx)=g*gr(ikx,iref)+f*gr(ikx,iref+1)
!                                ws(ikx)=gs(ikx,iref)
!                                wr(ikx)=gr(ikx,iref)
                            elseif (vel(iz,ikx) >= vref(nref,iz)) then
                                ws(ikx)=gs(ikx,nref)
                                wr(ikx)=gr(ikx,nref)
                            end if
                        else
                            ws(ikx)=gs(ikx,iref)
                            wr(ikx)=gr(ikx,iref)
                        end if                            
                    end do
                end do
           end if vxz
        end do depth
        if (mod(iw,10)==1) write(0,*)  "DONE FREQUENCY # ",iw," node = ",inode," OUT OF ",rec%w%n," / SHOT # ",ixs
end do omega

end subroutine owpspi

end module
