module OW_pspi

use OWvel
use OW_type
use OW_image
use OW_fftw
use wei_myfft_mod
use OW_grabvel
use OW_parms
use OW_imageIO

implicit none

public						:: owpspi_init, owpspi
private						:: isp,ixs

integer						:: isp,ixs
real						:: wb

contains

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

integer						:: isp_in,ixs_in
real						:: wb_in

isp=isp_in
ixs=ixs_in
wb=wb_in

end subroutine


!---------------------------------------------------------
subroutine owpspi(recv,source,taper,vel,image,vref1) 

real,optional					:: vref1(:,:,:)
integer						:: iw,iz,ikx,iref,iref1,ixw,ish,ino,i,j,iv,ipath
integer						:: nv,ihso,nwb,inode,stat,nrefvel
integer						:: indref(1)
real						:: sign, sqa,a,dwb,taper(:)
real						:: minv, maxv,vw,vpath
real, allocatable					:: GG(:,:,:,:),vref(:),f(:),g(:)
complex, allocatable				:: ws(:,:),wr(:,:),gs0(:,:),gr0(:,:),GGs(:,:,:),GGr(:,:,:)
complex,allocatable					:: tflds(:,:),tfldr(:,:),gs(:,:),gr(:,:)
complex						:: recv(:,:),source(:,:),phase
real						:: image(:,:,:),vel(:,:)
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,img%sh%n,img%z%n,node))
allocate(ws(kxm%n,node)); ws=0.
allocate(wr(kxm%n,node)); wr=0.
allocate(gs0(kxm%n,node));gs0=0
allocate(gr0(kxm%n,node));gr0=0
if (exist_file("ws").and.ixs.eq.1) then
   allocate(GGs(img%z%n,kxm%n,rec%w%n))
   allocate(GGr(img%z%n,kxm%n,rec%w%n))
end if
!--------------------------------------------------
nwb=floor(wb/img%z%d)+1
dwb=nwb*img%z%d
if (nwb.ne.1) then
   vw=minval(vel(nwb-1,:))
   write(0,*) "DATUMING WITH VELOCITY = ",vw
end if
!--------------------------------------------------
! Setting FFTW plans
allocate(tflds(kxm%n,node));tflds=0
call owfftws(tflds,node)

allocate(tfldr(kxm%n,node));tfldr=0
call owfftwr(tfldr,node)

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

  do iw = 1,rec%w%n-subf

        inode=omp_get_thread_num()+1; 
        ws(:,inode)=0;wr(:,inode)=0.
     if (marine) then
        ws(:,inode)=source(:,iw)
        wr(:,inode)=recv(:,iw)
     else
        ws(tapertr+padtraces+1:rec%h%n+padtraces+tapertr,inode)=source(:,iw)
        wr(tapertr+padtraces+1:rec%h%n+padtraces+tapertr,inode)=recv(:,iw)
     end if
        do iz=nwb+1,nzmax

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Compute vref
!
       if (present(vref1)) then
          nrefvel=count(mask = vref1(:,1,iz).ne.-1)
          if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
          vref=1/vref1(:,1,iz); 
          if (nrefvel==1) minv=vref(1)
       else
!          if (allocated(vref)) deallocate(vref); allocate(vref(nref))
!          call owvref(nref,vsal,vel,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

          nrefvel=nref
          if (maxval(vel).ge.vsal) then
             if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel+1))
             call owvref(nrefvel,vel(iz,:),vref(1:nrefvel),minv,maxv)
             vref=0;vref(nrefvel+1)=maxval(vel)
             nrefvel=nrefvel+1 ! added in 10/24/08
          else
             if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
             vref=0
             call owvref(nrefvel,vel(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));
          if (allocated(gr)) deallocate(gr); allocate(gr(kxm%n,nrefvel));
          gs(:,:)=0; gr(:,:)=0
       else
          gs0(:,inode)=0; gr0(:,inode)=0
       end if

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

            tfldr(:,inode)=wr(:,inode)
            call owfftexer(tfldr(:,inode),inode)
            wr(:,inode)=tfldr(:,inode)/sqrt(1.*kxm%n)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
!      Propagate wavefields when velocity is constant 
!
          if (nrefvel.eq.1) then
               gs0(:,inode)=ws(:,inode)
               gr0(:,inode)=wr(:,inode)
!
!      Operations in the wavenumber domain
!      phase shift
                do ikx = 1,kxm%n
                   sqa = rec%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)))
                       gr0(ikx,inode)=gr0(ikx,inode)*exp(cmplx(0,-v%z%d*sqrt(sqa)))
                   else
                       gs0(ikx,inode)=gs0(ikx,inode)*exp(-sqrt(-sqa)*v%z%d)
                       gr0(ikx,inode)=gr0(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)

                tfldr(:,inode)=gr0(:,inode)
                call owifftexer(tfldr(:,inode),inode)
                wr(:,inode)=tfldr(:,inode)/sqrt(1.*kxm%n)

            else
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!      Propagate wavefields with reference velocities
!
                 do iref=1,nrefvel

                    gs(:,iref)=ws(:,inode)
                    gr(:,iref)=wr(:,inode)

                    do ikx = 1,kxm%n
!      phase shift
                        sqa = rec%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)))
                            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

! Operations in x            

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

                    tfldr(:,inode)=gr(:,iref)
                    call owifftexer(tfldr(:,inode),inode)
                    gr(:,iref)=tfldr(:,inode)/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)));
                    gr(:,iref) = gr(:,iref) * exp(cmplx(0,-v%z%d*rec%w%all(iw))*(1./vel(iz,:)-1./vref(iref)));

                 end do
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
!       Interpolate between wavefields
!
!
                if (present(vref1)) then
                   do iref1=2,nrefvel
                      where (vel(iz,:) >= 1/vref1(iref1,2,iz) .and. vel(iz,:) < 1/vref1(iref1-1,2,iz))
                         ws(:,inode)=gs(:,iref1)
                         wr(:,inode)=gr(:,iref1)
                      end where
                   end do
                   where (vel(iz,:) > 1/vref1(1,2,iz))
                      ws(:,inode)=gs(:,1)
                      wr(:,inode)=gr(:,1)
                   end where
                else
                   do ikx=1,kxm%n
                       indref=0
                       indref=maxloc(vref, mask = int(vref).lt.int(vel(iz,ikx))) 
                       if (indref(1).ne.0.and.indref(1).ne.size(vref)-1) then
                          f(inode)=abs(vel(iz,ikx) - vref(indref(1)))/vel(iz,ikx)
                          f(inode)=f(inode)-int(f(inode))
                          g(inode)=1-f(inode)
!                          f(inode)=(vel(iz,ikx) - vref(indref(1)))/vel(iz,ikx)
!                          g(inode)=1-f(inode)
                          ws(ikx,inode)=g(inode)*gs(ikx,indref(1))+f(inode)*gs(ikx,indref(1)+1)
                          wr(ikx,inode)=g(inode)*gr(ikx,indref(1))+f(inode)*gr(ikx,indref(1)+1)
                       else
                          ws(ikx,inode)=gs(ikx,indref(1)+1)
                          wr(ikx,inode)=gr(ikx,indref(1)+1)
                       end if
                   end do
                end if  
            end if 
            ws(:,inode)=ws(:,inode)*taper
            wr(:,inode)=wr(:,inode)*taper
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Imaging Condition by crosscorrelation
!

           if ((iz-1)*img%z%d >= zmin) then
              phase = cmplx(cos(pi/4.), -sin(pi/4.)) * sqrt(rec%w%all(iw))
              do ish=1,img%sh%n
                  ihso=nint(((ish-1)*img%sh%d+img%sh%o)/(img%xm%d))
                  do ixw=1,kxm%n
                      if (ixw-ihso >= 1 .and. ixw-ihso <= kxm%n .and. ixw+ihso >= 1 .and. ixw+ihso <= kxm%n) then
                          GG(ixw,ish,iz-nzmin+1,inode)=GG(ixw,ish,iz-nzmin+1,inode) + &
                             real(conjg(ws(ixw-ihso,inode))*wr(ixw+ihso,inode)*phase)
                      end if
                  end do
              end do
           end if

           if (exist_file("ws").and.ixs.eq.1) then
              GGs(iz,:,iw)=ws(:,inode)
              GGr(iz,:,iw)=wr(:,inode)
           end if

        end do

        if (mod(iw,20)==1) write(0,5)  "DONE FREQUENCY # ",iw," OUT OF ",rec%w%n-subf," / SHOT # ",ixs
5       format(A17,i5,2(A8,i5))
end do
!$OMP END PARALLEL DO
write(0,*) "pspi 002"

if (exist_file("ws").and.ixs.eq.1) then
   call srite("ws",(GGs), 8*img%z%n*kxm%n*rec%w%n)
   deallocate(GGs)
end if
if (exist_file("wr").and.ixs.eq.1) then
   call srite("wr",GGr, 8*img%z%n*kxm%n*rec%w%n)
   deallocate(GGr)
end if

call owfftdestroy()
deallocate(tfldr,tflds)

   do i=1,node
      !$OMP parallel do private (ixw,iz,ish)
      do iz=1,img%z%n
         do ish=1,img%sh%n
            do ixw=1,kxm%n
               image(ixw,ish,iz)=image(ixw,ish,iz)+GG(ixw,ish,iz,i)
            end do
         end do
      end do
      !$OMP END PARALLEL DO
   end do

deallocate(GG,ws,wr,gs0,gr0)

end subroutine
end module

