module OW_pspi_areal

use OW_type
use OW_parms_areal
use OW_fftw_areal
use OWvelo_areal

implicit none

integer, private :: iz
contains

!---------------------------------------------------------
subroutine owpspi_init(iz_in)
integer :: iz_in
iz=iz_in
end subroutine 
!---------------------------------------------------------
subroutine owpspi(iw,inode,sign,ws,vel,tfld,gat,vref1) 

real,optional					:: vref1(:,:)
character*3					:: gat
integer						:: iw,ikx,iref,iref1,ixw,i,j
integer						:: nwb,inode,nrefvel,sign
integer						:: indref(1)
real						:: sqa,a,dwb,f,g,sqav(kxm%n)
real						:: minv, maxv, dvel
real, allocatable					:: vref(:)
complex						:: ws(:)
complex, allocatable				:: gs(:,:)
complex						:: tfld(:)
real						:: vel(:),vel1(kxm%n)
integer, external 					:: omp_get_thread_num,omp_get_num_threads

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Compute vref
!
       vel1=1/vel
       if (present(vref1)) then
          nrefvel=count(mask = vref1(:,1).ne.-1)
          if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
          vref=1/vref1(:,1); 
          if (nrefvel==1) minv=vref(1)
       else
          if (allocated(vref)) deallocate(vref); allocate(vref(nref))
          call owvref(vel1,vref,minv,maxv,dvel)
          if (int(minv).eq.int(maxv)) then
             nrefvel=1
             if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
             vref=maxv
          else
             nrefvel=nref
          end if
       end if

       if (allocated(gs)) deallocate(gs); allocate(gs(kxm%n,nrefvel));gs=0
       if (gat.eq."sou") then
          tfld=ws
          call owfftexes(tfld,inode)
          ws=tfld/sqrt(1.*kxm%n)
       else
          tfld=ws
          call owfftexer(tfld,inode)
          ws=tfld/sqrt(1.*kxm%n)
       end if

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!      Propagate wavefields with reference velocities
!
       do iref=1,nrefvel
          gs(:,iref)=ws
          sqav=(rec%w%all(iw) / vref(iref))**2 - kxm%all**2
          where(sqav >= 0) 
              gs(:,iref)=gs(:,iref)*exp(cmplx(0,sign*v%z%d*sqrt(sqav)))
          elsewhere
              gs(:,iref)=gs(:,iref)*exp(-sqrt(-sqav)*v%z%d)
          end where
          if (gat.eq."sou") then
             tfld=gs(:,iref)
             call owifftexes(tfld,inode)
             gs(:,iref)=tfld/sqrt(1.*kxm%n)
          else
             tfld=gs(:,iref)
             call owifftexer(tfld,inode)
             gs(:,iref)=tfld/sqrt(1.*kxm%n)
          end if
!      split step
          gs(:,iref) = gs(:,iref) * exp(cmplx(0,sign*v%z%d*rec%w%all(iw))*(vel-1./vref(iref)));
       end do
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!Interpolate wavefields
!
       if (nrefvel.eq.1) then
          ws=gs(:,1)
       else
!                  if (iw.eq.1.and.gat.eq."sou") write(0,*) iz,dvel
          if (present(vref1)) then
             do iref1=2,nrefvel
                where (vel1 >= 1/vref1(iref1,2) .and. vel1 < 1/vref1(iref1-1,2))
                   ws=gs(:,iref1)
                end where
             end do
             where (vel1 > 1/vref1(1,2))
                ws=gs(:,1)
             end where
          else
             do ikx=1,kxm%n
                 indref=0
                 indref=maxloc(vref, mask = int(vref).lt.int(vel1(ikx))) 
                 if (indref(1).ne.0.and.indref(1).ne.size(vref)-1) then
                    f=abs(vel1(ikx) - vref(indref(1)))/vel1(ikx)
                    f=f-int(f)
                    g=1-f
                    ws(ikx)=g*gs(ikx,indref(1))+f*gs(ikx,indref(1)+1)
                 else
                    ws(ikx)=gs(ikx,indref(1)+1)
                 end if
             end do
!             do ikx=1,kxm%n
!                 indref=0
!                 indref=maxloc(vref, mask = int(vref).lt.int(vel1(ikx)))
!                 indref=indref+1 
!!                 if (indref(1).ne.0.and.indref(1).ne.size(vref)-1) then
!!                    f=abs(vel1(ikx) - vref(indref(1)))/vel1(ikx)
!!                    f=f-int(f)
!!                    f=(vref(indref(1)) - vel1(ikx))/(vref(indref(1)) - vref(indref(1)+1))
!                    f=(vref(indref(1))-vel1(ikx))/dvel
!                    g=1-f
!!                 if (iw.eq.1.and.gat.eq."sou") write(13,13) iz,ikx,indref,int(vel1(ikx)),int(vref),f,g
!!13               format(3i4,5I5,2f6.3)
!
!                 if (indref(1).ne.nrefvel) then
!                    ws(ikx)=f*gs(ikx,indref(1))+g*gs(ikx,indref(1)+1)
!                 else
!                    ws(ikx)=gs(ikx,indref(1))
!                 end if
!
!!                 if (indref(1).ne.1) then
!!                    ws(ikx)=f*gs(ikx,indref(1)-1)+g*gs(ikx,indref(1))
!!!                    ws(ikx)=g*gs(ikx,indref(1)-1)+f*gs(ikx,indref(1))
!!                 else
!!                    ws(ikx)=gs(ikx,indref(1))
!!                 end if
!
!             end do
          end if  
       end if

       if (allocated(vref)) deallocate(vref) 
       if (allocated(gs)) deallocate(gs)

end subroutine owpspi

!======================================================================================
end module
