module OW_pspi

use OWvel
use OW_type
use OW_fftw
use OW_parms

implicit none

public						:: owpspi_init, owpspi

integer,private					:: ixs,ni
real						:: wb

contains

!---------------------------------------------------------
subroutine owpspi_init(ixs_in,ni_in,wb_in)

integer						:: ixs_in,ni_in
real						:: wb_in

ni=ni_in
ixs=ixs_in
wb=wb_in

end subroutine


!---------------------------------------------------------
subroutine owpspi(source_in,taper_in,vel_in,refl_in,vref_in,illum_in)

real,optional					:: vref_in(:,:,:),illum_in(:,:)
integer						:: iw,iz,ikx,iref,iref1,ixw,ish,ino,i,j
integer						:: nv,ihso,nwb,inode,stat,nrefvel
integer						:: indref(1),ikxx,iks
real						:: sign, sqa,a,dwb,taper_in(:),refl_in(:,:)
real						:: minv, maxv,vw,vpath
real, allocatable					:: vref(:),f(:),g(:)
complex, allocatable				:: ws(:,:),gs0(:,:),ss(:,:,:)
complex,allocatable					:: GG(:,:,:),tflds(:,:),gs(:,:),GGG(:,:,:),shot(:,:)
complex						:: source_in(:,:),phase
real						:: vel_in(:,:)
integer, external 					:: omp_get_thread_num,omp_get_num_threads

if (node.eq.9999) then
  !$OMP PARALLEL
  node = omp_get_num_threads()
  !$OMP END PARALLEL
write(0,*) " NUM.THREADS = ",node
end if

allocate(f(node),g(node));f=0;g=0
allocate(GG(kxm%n,sou%w%n,img%z%n))    
allocate(ss(kxm%n,node,img%z%n));ss=0
allocate(ws(kxm%n,node)); ws=0.
allocate(gs0(kxm%n,node));gs0=0
allocate(GGG(kxm%n,sou%w%n,img%z%n));GGG=0.

!--------------------------------------------------
nwb=floor(wb/img%z%d)+1
dwb=nwb*img%z%d
if (nwb.ne.1) then
   vw=minval(vel_in(nwb-1,:))
   write(0,*) "DATUMING WITH VELOCITY = ",vw
end if
!--------------------------------------------------
! Setting FFTW plans
allocate(tflds(kxm%n,node));tflds=0
call owfftws(tflds,node)

!$OMP parallel do private (inode,&
                             iw,&
                             iz,&
                            ikx,&
                          iref1,&
                         indref,&
                           iref,&
                            ish,&
                           ihso,&
                            ixw,&
                            sqa,&
                           minv,&
                           vref,&
                             gs,&
                        nrefvel,&
                           maxv,&
                           phase) schedule(dynamic,1) num_threads(node)

  do iw = 1,sou%w%n

        inode=omp_get_thread_num()+1; 
        ws(:,inode)=0
        ws(:,inode)=source_in(:,iw)
        GG(:,iw,1)=ws(:,inode)
        do iz=2,img%z%n

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Compute vref
!
       if (present(vref_in)) then
          nrefvel=count(mask = vref_in(:,1,iz).ne.-1)
          if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
          vref=1/vref_in(:,1,iz); 
          if (nrefvel==1) minv=vref(1)
       else
          nrefvel=nref
          if (maxval(vel_in).ge.vsal) then
             if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel+1))
             call owvref(nrefvel,vel_in(iz,:),vref(1:nrefvel),minv,maxv)
             vref=0;vref(nrefvel+1)=maxval(vel_in(iz,:))
             nrefvel=nrefvel+1 ! added in 10/24/08
          else
             if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
             vref=0
             call owvref(nrefvel,vel_in(iz,:),vref,minv,maxv)
             if (minv.eq.maxv) then
                nrefvel=1
                if (allocated(vref)) then
                   deallocate(vref)
                   allocate(vref(nrefvel))
                   vref=maxv
                end if
             end if
          end if
       end if

       if (nrefvel.gt.1) then
          if (allocated(gs)) deallocate(gs); allocate(gs(kxm%n,nrefvel));
          gs(:,:)=0
       else
          gs0(:,inode)=0
       end if

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!  
            tflds(:,inode)=ws(:,inode)
            call owfftexes(tflds(:,inode),inode)
            ws(:,inode)=tflds(:,inode)/sqrt(1.*kxm%n)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
!      Propagate wavefields when velocity is constant 
!
          if (nrefvel.eq.1) then
               gs0(:,inode)=ws(:,inode)
!
!      Operations in the wavenumber domain
!      phase shift
                do ikx = 1,kxm%n
                   sqa = sou%w%all(iw)**2 /minv**2 - kxm%all(ikx)**2
                   if (sqa >= 0) then
                       gs0(ikx,inode)=gs0(ikx,inode)*exp(cmplx(0,v%z%d*sqrt(sqa)))
                   else
                       gs0(ikx,inode)=gs0(ikx,inode)*exp(-sqrt(-sqa)*v%z%d)
                   end if
                end do
                tflds(:,inode)=gs0(:,inode)
                call owifftexes(tflds(:,inode),inode)
                ws(:,inode)=tflds(:,inode)/sqrt(1.*kxm%n)
          else
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!      Propagate wavefields with reference velocities
!
                 do iref=1,nrefvel
                    gs(:,iref)=ws(:,inode)
                    do ikx = 1,kxm%n
!      phase shift
                        sqa = sou%w%all(iw)**2 / vref(iref)**2 - kxm%all(ikx)**2
                        if (sqa > 0) then
                            gs(ikx,iref)=gs(ikx,iref)*exp(cmplx(0,v%z%d*sqrt(sqa)))
                        else
                            gs(ikx,iref)=gs(ikx,iref)*exp(-sqrt(-sqa)*v%z%d)
                        end if
                    end do
! Operations in x            
                    tflds(:,inode)=gs(:,iref)
                    call owifftexes(tflds(:,inode),inode)
                    gs(:,iref)=tflds(:,inode)/sqrt(1.*kxm%n)
!      split step
                    gs(:,iref) = gs(:,iref) * exp(cmplx(0,v%z%d*sou%w%all(iw))*(1./vel_in(iz,:)-1./vref(iref)));
                 end do

!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
!       Interpolate wavefields
!
!
                if (present(vref_in)) then
                   do iref1=2,nrefvel
                      where (vel_in(iz,:) >= 1/vref_in(iref1,2,iz) .and. vel_in(iz,:) < 1/vref_in(iref1-1,2,iz))
                         ws(:,inode)=gs(:,iref1)
                      end where
                   end do
                   where (vel_in(iz,:) > 1/vref_in(1,2,iz))
                      ws(:,inode)=gs(:,1)
                   end where
                else
                   do ikx=1,kxm%n
                       indref=0
                       indref=maxloc(vref, mask = int(vref).lt.int(vel_in(iz,ikx))) 
                       if (indref(1).ne.0.and.indref(1).ne.size(vref)-1) then
                          f(inode)=abs(vel_in(iz,ikx) - vref(indref(1)))/vel_in(iz,ikx)
                          f(inode)=f(inode)-int(f(inode))
                          g(inode)=1-f(inode)
!                          f(inode)=(vel_in(iz,ikx) - vref(indref(1)))/vel_in(iz,ikx)
!                          g(inode)=1-f(inode)
                          ws(ikx,inode)=g(inode)*gs(ikx,indref(1))+f(inode)*gs(ikx,indref(1)+1)
                       else
                          ws(ikx,inode)=gs(ikx,indref(1)+1)
                       end if
                   end do
                end if  
          end if 
          ws(:,inode)=ws(:,inode)*taper_in
          GG(:,iw,iz)=ws(:,inode)
        end do
        if (light) then
           do iz=1,img%z%n
              illum_in(:,iz)=illum_in(:,iz)+GG(:,iw,iz)*conjg(GG(:,iw,iz))
           end do
        end if
!!!!!!!
!
! Scattered wavefield
!
        ws(:,inode)=0.
        ss(:,inode,:)=0.

        do iz=1,img%z%n
           ss(:,inode,iz)=ss(:,inode,iz)+GG(:,iw,iz)*refl_in(iz,:)
        end do

        do iz=img%z%n,1,-1

            ws(:,inode)=ws(:,inode)+ss(:,inode,iz)
            ws(:,inode)=ws(:,inode)*taper_in

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

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Compute vref
!
       if (present(vref_in)) then
          nrefvel=count(mask = vref_in(:,1,iz).ne.-1)
          if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
          vref=1/vref_in(:,1,iz); 
          if (nrefvel==1) minv=vref(1)
       else
          nrefvel=nref
          if (maxval(vel_in).ge.vsal) then
             if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel+1))
             call owvref(nrefvel,vel_in(iz,:),vref(1:nrefvel),minv,maxv)
             vref=0;vref(nrefvel+1)=maxval(vel_in(iz,:))
             nrefvel=nrefvel+1 ! added in 10/24/08
          else
             if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
             vref=0
             call owvref(nrefvel,vel_in(iz,:),vref,minv,maxv)
             if (minv.eq.maxv) then
                nrefvel=1
                if (allocated(vref)) then
                   deallocate(vref)
                   allocate(vref(nrefvel))
                   vref=maxv
                end if
             end if
          end if
       end if
       if (nrefvel.gt.1) then
          if (allocated(gs)) deallocate(gs); allocate(gs(kxm%n,nrefvel));
          gs(:,:)=0
       else
          gs0(:,inode)=0
       end if


!      Propagate wavefields when velocity is constant 
!
          if (nrefvel.eq.1) then
               gs0(:,inode)=ws(:,inode)
!
!      Operations in the wavenumber domain
!      phase shift
                do ikx = 1,kxm%n
                   sqa = sou%w%all(iw)**2 /minv**2 - kxm%all(ikx)**2
                   if (sqa > 0) then
                       gs0(ikx,inode)=gs0(ikx,inode)*exp(cmplx(0,v%z%d*sqrt(sqa)))
                   else
                       gs0(ikx,inode)=gs0(ikx,inode)*exp(-sqrt(-sqa)*v%z%d)
                   end if
                end do
                tflds(:,inode)=gs0(:,inode)
                call owifftexes(tflds(:,inode),inode)
                ws(:,inode)=tflds(:,inode)/sqrt(1.*kxm%n)
          else
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!      Propagate wavefields with reference velocities
!
                 do iref=1,nrefvel
                    gs(:,iref)=ws(:,inode)
                    do ikx = 1,kxm%n
!      phase shift
                        sqa = sou%w%all(iw)**2 / vref(iref)**2 - kxm%all(ikx)**2
                        if (sqa > 0) then
                            gs(ikx,iref)=gs(ikx,iref)*exp(cmplx(0,v%z%d*sqrt(sqa)))
                        else
                            gs(ikx,iref)=gs(ikx,iref)*exp(-sqrt(-sqa)*v%z%d)
                        end if
                    end do
! Operations in x            
                    tflds(:,inode)=gs(:,iref)
                    call owifftexes(tflds(:,inode),inode)
                    gs(:,iref)=tflds(:,inode)/sqrt(1.*kxm%n)
!      split step
                    gs(:,iref) = gs(:,iref) * exp(cmplx(0,v%z%d*sou%w%all(iw))*(1./vel_in(iz,:)-1./vref(iref)));
                 end do
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
!       Interpolate between wavefields
!
!
                if (present(vref_in)) then
                   do iref1=2,nrefvel
                      where (vel_in(iz,:) >= 1/vref_in(iref1,2,iz) .and. vel_in(iz,:) < 1/vref_in(iref1-1,2,iz))
                         ws(:,inode)=gs(:,iref1)
                      end where
                   end do
                   where (vel_in(iz,:) > 1/vref_in(1,2,iz))
                      ws(:,inode)=gs(:,1)
                   end where
                else
                   do ikx=1,kxm%n
                       indref=0
                       indref=maxloc(vref, mask = int(vref).lt.int(vel_in(iz,ikx))) 
                       if (indref(1).ne.0.and.indref(1).ne.size(vref)-1) then
                          f(inode)=abs(vel_in(iz,ikx) - vref(indref(1)))/vel_in(iz,ikx)
                          f(inode)=f(inode)-int(f(inode))
                          g(inode)=1-f(inode)
!                          f(inode)=(vel_in(iz,ikx) - vref(indref(1)))/vel_in(iz,ikx)
!                          g(inode)=1-f(inode)
                          ws(ikx,inode)=g(inode)*gs(ikx,indref(1))+f(inode)*gs(ikx,indref(1)+1)
                       else
                          ws(ikx,inode)=gs(ikx,indref(1)+1)
                       end if
                   end do
                end if  
          end if 
          GGG(:,iw,iz)=ws(:,inode)
        end do

        if (mod(iw,50)==1) write(0,5)  "DONE FREQUENCY # ",iw," OUT OF ",sou%w%n," / Gr.F  # ",ixs
5       format(A17,i5,2(A8,i5))

end do
!$OMP END PARALLEL DO
allocate(shot(kxm%n,sou%w%n));shot=0.
do iw=1,sou%w%n
!   phase = cmplx(cos(pi/4.), -sin(pi/4.)) * sqrt(sou%w%all(iw))
   shot(:,iw)=GGG(:,iw,1)!*phase
end do
if (exist_file("GG")) then
   call srite("GG",GG, 8*sou%w%n*kxm%n*img%z%n)
end if
if (exist_file("GGG")) then
   call srite("GGG",GGG, 8*sou%w%n*kxm%n*img%z%n)
end if
call srite("shots",shot, 8*sou%w%n*kxm%n)
deallocate(GGG,shot,ss)
call owfftdestroy()
deallocate(tflds)
deallocate(GG,ws,gs0)
end subroutine
end module

