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,vel,image) 

integer						:: iw,iz,ikx,iref,ixw,ish,ino,i
integer						:: nv,ihso,nwb,inode,stat
integer						:: indref(1)
real						:: sign, sqa,a,dwb
real						:: minv, maxv,vw
real, allocatable					:: vref(:),f(:),g(:),GG(:,:,:,:)
complex, allocatable				:: ws(:,:),wr(:,:),gs0(:,:),gr0(:,:)
complex,allocatable					:: tflds(:,:),tfldr(:,:),gs(:,:,:),gr(:,:,:)
complex						:: recv(:,:),source(:,:)
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(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(gs(kxm%n,nref+1,node)); gs=0
allocate(gr(kxm%n,nref+1,node)); gr=0
allocate(gs0(kxm%n,node));gs0=0
allocate(gr0(kxm%n,node));gr0=0
allocate(f(node));f=0
allocate(g(node));g=0

!--------------------------------------------------
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,&
                           iref,&
                            ish,&
                           ihso,&
                         indref,&
                            ixw,&
                            sqa,&
                           minv,&
                           vref,&
                            maxv) schedule(dynamic,1) num_threads(node)

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

        inode=omp_get_thread_num()+1; 

        ws(:,inode)=source(:,iw)
        wr(:,inode)=recv(:,iw) 


        if (nwb.ne.1) then

           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)

           do ikx = 1,kxm%n
              sqa = rec%w%all(iw)**2 / vw**2 - kxm%all(ikx)**2
              if (sqa > 0) then
                 wr(ikx,inode)=wr(ikx,inode)*exp(cmplx(0,dwb*sqrt(sqa)))
                 ws(ikx,inode)=ws(ikx,inode)*exp(cmplx(0,-dwb*sqrt(sqa)))
              else
                 wr(ikx,inode)=wr(ikx,inode)*exp(-sqrt(-sqa)*dwb)
                 ws(ikx,inode)=ws(ikx,inode)*exp(-sqrt(-sqa)*dwb)
              end if
           end do

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

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

!           wr(:,inode)=wr(:,inode)*taper
!           ws(:,inode)=ws(:,inode)*taper
        end if

        depth:do iz=nwb+1,img%z%n

!!--------------------------------------------------
!! Compute vref
!!
            if (maxval(vel(iz,:)).ge.vsal) then
               if (allocated(vref)) deallocate(vref); allocate(vref(nref+1))
               vref=0; gs(:,:,inode)=0; gr(:,:,inode)=0
               call owvref(nref,vel(iz,:),vref(1:nref),minv,maxv)
               vref(nref+1)=maxval(vel(iz,:))
            else if(maxval(vel(iz,:)).lt.vsal) then
               if (allocated(vref)) deallocate(vref); allocate(vref(nref))
               vref=0; gs(:,:,inode)=0; gr(:,:,inode)=0
               call owvref(nref,vel(iz,:),vref(1:nref),minv,maxv)
            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)

           vxz:if (minv .eq. maxv .and. maxv.eq.vref(size(vref))) then
               gs0(:,inode)=ws(:,inode)
               gr0(:,inode)=wr(:,inode)
!
!      Operations in the wavenumber domain
!      phase shift
                wavenumber1:do ikx = 1,kxm%n
                   sqa = rec%w%all(iw)**2 / minv**2 - kxm%all(ikx)**2

                   evanescent1: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 evanescent1
                end do wavenumber1

                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
                velref: do iref=1,size(vref)                 
                    gs(:,iref,inode)=ws(:,inode)
                    gr(:,iref,inode)=wr(:,inode); 
                    wavenumber2: do ikx = 1,kxm%n
!      phase shift
                        sqa = rec%w%all(iw)**2 / vref(iref)**2 - kxm%all(ikx)**2;
                        evanescent2: if (sqa > 0) then
                            gs(ikx,iref,inode)=gs(ikx,iref,inode)*exp(cmplx(0,v%z%d*sqrt(sqa)));
                            gr(ikx,iref,inode)=gr(ikx,iref,inode)*exp(cmplx(0,-v%z%d*sqrt(sqa)));
                        else
                            gs(ikx,iref,inode)=gs(ikx,iref,inode)*exp(-sqrt(-sqa)*v%z%d)
                            gr(ikx,iref,inode)=gr(ikx,iref,inode)*exp(-sqrt(-sqa)*v%z%d)
                        end if evanescent2
                    end do wavenumber2

! Operations in x            

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

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

!      split step
                    gs(:,iref,inode) = gs(:,iref,inode) * exp(cmplx(0,v%z%d*rec%w%all(iw))*(1./vel(iz,:)-1./vref(iref)));
                    gr(:,iref,inode) = gr(:,iref,inode) * exp(cmplx(0,-v%z%d*rec%w%all(iw))*(1./vel(iz,:)-1./vref(iref)));
                end do velref

!                do iref=1,nref-1
!                interpwvfld: where (vel(iz,:) >= vref(iref) .and. vel(iz,:) < vref(iref+1))
!                     ws=gs(:,iref)
!                     wr=gr(:,iref)
!                     end where interpwvfld
!                end do
!                where (vel(iz,:) >= vref(nref))
!                      ws=gs(:,nref)
!                      wr=gr(:,nref)
!                end where

                do ikx=1,kxm%n
                    indref=0
                    indref=maxloc(vref, mask = vref.lt.vel(iz,ikx)) 
                    if (indref(1).ne.0.and.indref(1).ne.size(vref)-1) then
                       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),inode)+f(inode)*gs(ikx,indref(1)+1,inode)
                       wr(ikx,inode)=g(inode)*gr(ikx,indref(1),inode)+f(inode)*gr(ikx,indref(1)+1,inode)
                    else
                       ws(ikx,inode)=gs(ikx,indref(1)+1,inode)
                       wr(ikx,inode)=gr(ikx,indref(1)+1,inode)
                    end if
                end do
           end if vxz

!
! Imaging Condition by crosscorrelation
!
           do ish=1,img%sh%n
               ihso=floor(((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,inode)=GG(ixw,ish,iz,inode)+real(conjg(ws(ixw-ihso,inode))*(wr(ixw+ihso,inode)))
                   end if
               end do
           end do

        end do depth

        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 omega
!$OMP END PARALLEL DO

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

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

end subroutine

end module

