module OW_pspi_areal

use OW_type
use OW_fftw_areal
use OW_parms_areal

implicit node

contains

subroutine owpspi(ws,sref,vel,sign,iwa,GGs)

complex,optional					:: GGs(:,:,:)
integer						:: iwa,iz,ikx,iref,nrefvel,sign
real						:: sqa,vel(:,:),sref(:,:,:)
complex						:: ws(:)
complex,allocatable					:: tflds(:),gs0(:),gs(:,:)

!
! sign = 1 for downward extrapolating shots
! sign = -1 for downward extrapolating receivers
!
 
allocate(gs0(kxm%n))
gs0=0

!--------------------------------------------------
! Compute vref
!
   nrefvel=count(mask = sref(:,1,iz).ne.-1)
      
   if (nrefvel.gt.1) then
      if (allocated(gs)) deallocate(gs); allocate(gs(kxm%n,nrefvel));
      gs=0
   else
      gs0=0
   end if

   if (nrefvel .eq. 1) then
      gs0=ws
      do ikx = 1,kxm%n
         sqa = (rec%w%all(iwa) * sref(1,1,iz))**2 - kxm%all(ikx)**2
         if (sqa > 0) then
            gs0(ikx)=gs0(ikx)*exp(cmplx(0,sign*v%z%d*sqrt(sqa)))
         else
            gs0(ikx)=gs0(ikx)*exp(-sqrt(-sqa)*v%z%d)
         end if
      end do

      tflds=gs0
      call owifftexes(tflds,inode)
      ws=tflds/sqrt(1.*kxm%n)

   else

      do iref=1,nrefvel                 
         gs(:,iref)=ws
         do ikx = 1,kxm%n
            sqa = (rec%w%all(iwa) * sref(iref,1,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)));
            else
               gs(ikx,iref)=gs(ikx,iref)*exp(-sqrt(-sqa)*v%z%d)
            end if
         end do

         tflds=gs(:,iref)
         call owifftexes(tflds,inode)
         gs(:,iref)=tflds/sqrt(1.*kxm%n)

         gs(:,iref) = gs(:,iref) * exp(cmplx(0,sign*v%z%d*rec%w%all(iwa))*(1./vel(iz,:)-sref(iref,1,iz)));

      end do

      do iref=1,nrefvel-1
         where (vel(iz,:) >= 1/sref(iref,1,iz) .and. vel(iz,:) < 1/sref(iref+1,1,iz))
            ws=gs(:,iref)
         end where
      end do
      where (vel(iz,:) >= 1/sref(nrefvel,1,iz))
         ws=gs(:,nrefvel)
      end where 
 
   end if

end subroutine

end module
