module OW_1pspi

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

private						:: isp,ixs,ni

integer						:: isp,ixs,ni
real						:: wb

contains

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

integer						:: isp_in,ni_in,ixs_in
real						:: wb_in

isp=isp_in
ixs=ixs_in
wb=wb_in
ni=ni_in
end subroutine
!---------------------------------------------------------
subroutine ow1pspi(recv,source,vel) 

integer						:: iw,iz,ikx,iref,ixw,ish,ino
integer						:: nv,ihso,nwb,node,inode,stat
integer,allocatable					:: indref(:)
real						:: sign, sqa,f,g
real						:: minv, maxv,vel(:,:)
real, allocatable					:: vref(:),sqa0(:)
complex, allocatable				:: ws(:),wr(:),gs0(:),gr0(:),gs(:,:),gr(:,:),tfld(:),GSS(:,:),GRR(:,:)
complex						:: recv(:,:),source(:)

integer, external 					:: omp_get_thread_num,omp_get_num_threads,omp_get_num_procs
!integer, external 					:: omp_set_num_threads

real,allocatable					:: image(:,:,:),GG(:,:,:)


stat=1
sign=-1;
node=omp_get_num_threads()
node=8

allocate(GG(kxm%n,img%sh%n,img%z%n))
allocate(image(kxm%n,img%sh%n,img%z%n))

!--------------------------------------------------
nwb=floor(wb/img%z%d)+1

!--------------------------------------------------
! Setting FFTW plans
allocate(tfld(kxm%n));tfld=0
call owfftw(tfld)
deallocate(tfld)

!$OMP parallel do private (inode,&
                             iw,&
                             iz,&
                            ikx,&
                           iref,&
                            ish,&
                           ihso,&
                            ixw,&
                             wr,&
                             gr,&
                            gr0,&
                             ws,&
                             gs,&
                            gs0,&
                           vref,&
                           sqa0,&
                            sqa,&
                           minv,&
                           maxv,&
                           tfld,&
                         indref,&
                              f,&
                              g) schedule(dynamic,1) num_threads(1)
! !$OMP DO schedule(dynamic,1) 

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

        if(allocated(ws)) deallocate(ws); allocate(ws(kxm%n)); ws=0.
        if(allocated(wr)) deallocate(wr); allocate(wr(kxm%n)); wr=0.

        inode=omp_get_thread_num()+1; 

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

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

            if(allocated(tfld)) deallocate(tfld); allocate(tfld(kxm%n));tfld=0

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


!!---------------------------------------------------
!!
!            tfld=ws; call sfftw_execute(planF); ws=tfld/sqrt(1.*kxm%n);
!            tfld=wr; call sfftw_execute(planF); wr=tfld/sqrt(1.*kxm%n)

            tfld=ws; call FFT1DX(tfld,1); ws=tfld;
            tfld=wr; call FFT1DX(tfld,1); wr=tfld;

           vxz:if (minv .eq. maxv .and. maxv.eq.vref(size(vref))) then
               if(allocated(gs0)) deallocate(gs0); allocate(gs0(kxm%n));gs0=ws
               if(allocated(gr0)) deallocate(gr0); allocate(gr0(kxm%n));gr0=wr
!
!      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)=gs0(ikx)*exp(cmplx(0,v%z%d*sqrt(sqa)))
                       gr0(ikx)=gr0(ikx)*exp(cmplx(0,-v%z%d*sqrt(sqa)))
                   else
                       gs0(ikx)=gs0(ikx)*exp(-sqrt(-sqa)*v%z%d)
                       gr0(ikx)=gr0(ikx)*exp(-sqrt(-sqa)*v%z%d)
                   end if evanescent1
                end do wavenumber1

            tfld=gs0; call FFT1DX(tfld,-1); ws=tfld;
            tfld=gr0; call FFT1DX(tfld,-1); wr=tfld;

!             tfld=gs0; call sfftw_execute(planI); ws=tfld/sqrt(1.*kxm%n)
!             tfld=gr0; call sfftw_execute(planI); wr=tfld/sqrt(1.*kxm%n)
!
            else
               if(allocated(sqa0)) deallocate(sqa0); allocate(sqa0(kxm%n));sqa0=0.
                velref: do iref=1,size(vref)                 
                    gs(:,iref)=ws;  gr(:,iref)=wr; 
                    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)=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 evanescent2
                    end do wavenumber2

! Operations in x            

            tfld=gs(:,iref); call FFT1DX(tfld,-1); gs(:,iref)=tfld;
            tfld=gr(:,iref); call FFT1DX(tfld,-1); gr(:,iref)=tfld;

!             tfld=gs(:,iref); call sfftw_execute(planI); gs(:,iref)=tfld/sqrt(1.*kxm%n)
!             tfld=gr(:,iref); call sfftw_execute(planI); gr(:,iref)=tfld/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 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
                    if(allocated(indref)) deallocate(indref); allocate(indref(1));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=(vel(iz,ikx) - vref(indref(1)))/vel(iz,ikx)
                       g=1-f
                       ws(ikx)=g*gs(ikx,indref(1))+f*gs(ikx,indref(1)+1)
                       wr(ikx)=g*gr(ikx,indref(1))+f*gr(ikx,indref(1)+1)
                    else
                       ws(ikx)=gs(ikx,indref(1)+1)
                       wr(ikx)=gr(ikx,indref(1)+1)
                    end if
                end do
           end if vxz

! Imaging Condition
!
           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)=GG(ixw,ish,iz)+real(conjg(ws(ixw-ihso))*(wr(ixw+ihso)))             ! correlation to image
                   end if
               end do
           end do

        end do depth
!
        !$OMP CRITICAL
           image=image+GG
           call owimage_io(image,ni,ixs)
        !$OMP END CRITICAL
          if (mod(iw,20)==1) write(0,5)  "DONE FREQUENCY # ",iw," OUT OF ",rec%w%n," / SHOT # ",ixs
5                            format(A17,i5,2(A8,i5))


end do omega

!$OMP END PARALLEL DO

stat=0

end subroutine

end module
