module OW_pspi2

use OW_type
use OW_param
use OW_fftw

implicit none

public						:: nrefer,isp1,iwa,ixs1
integer						:: nrefer,isp1,iwa,ixs1

contains

!-----------------------------------------------------------------
subroutine owpspi_init(iwa_in,nrefer_in,isp1_in,ixs1_in)

integer						:: ixs1_in,nrefer_in,isp1_in,iwa_in

iwa=iwa_in
nrefer=nrefer_in
isp1=isp1_in
ixs1=ixs1_in

end subroutine

!-------------------------------------------------------------------
subroutine owpspi(vel,vref,recv,source,GG)

integer						:: iz,ikx,iref,ixw,ish
integer						:: ihso
real						:: sqa
real						:: minv, maxv
real 						:: vref(:,:,:),vel(:,:)
real, allocatable				:: sqa0(:)
complex, allocatable				:: tfield(:)
complex, allocatable				:: ws(:),wr(:),gs0(:),gr0(:),gs(:,:),gr(:,:)
complex						:: recv(:,:),source(:,:,:),GG(:,:,:)

allocate(ws(rec%h%n));		ws=0
allocate(wr(rec%h%n));		wr=0
allocate(tfield(rec%h%n));	tfield=0
allocate(gs(rec%h%n,nrefer));	gs=0
allocate(gr(rec%h%n,nrefer));	gr=0
allocate(gs0(rec%h%n));		gs0=0
allocate(gr0(rec%h%n));		gr0=0
allocate(sqa0(rec%h%n));	sqa0=0

        ws(isp1)=source(1,iwa,1) ! to be modified to handle generalized sources
        wr=recv(:,iwa)
!        if (mod(iwa,10)==1) write(2,*)  "1 ",iwa,maxval(cabs(ws)),maxval(cabs(wr))

        depth:do iz=2,img%z%n

            tfield=ws
            call sfftw_execute(planF)
            ws=tfield/sqrt(1.*kxm%n)
            tfield=wr
            call sfftw_execute(planF)
            wr=tfield/sqrt(1.*kxm%n)
            minv=minval(vel(iz,:)); maxv=maxval(vel(iz,:))

            vxz:if (minv .eq. maxv) then
                gs0=ws;  gr0=wr;
!      Operations in the wavenumber domain
!                wavenumber1:do ik = 1,kxm%n
!      phase shift
!                   sqa = rec%w%all(iwa)**2 / minv**2 - kxm%all(ik)**2
!                   evanescent1:if (sqa > 0) then
!                       gs0(ik)=gs0(ik)*exp(cmplx(0,v%z%d*sqrt(sqa)))
!                       gr0(ik)=gr0(ik)*exp(cmplx(0,-v%z%d*sqrt(sqa)))
!                   else
!                       gs0(ik)=gs0(ik)*exp(-sqrt(-sqa)*v%z%d)
!                       gr0(ik)=gr0(ik)*exp(-sqrt(-sqa)*v%z%d)
!                   end if evanescent1
!                end do wavenumber1
        if (mod(iwa,10)==1) write(2,*)  "2 ",iwa,iz,maxval(cabs(gs0)),maxval(cabs(gr0))
                   sqa0 = rec%w%all(iwa)**2 / minv**2 - kxm%all**2
                   evanescent1:where (sqa0 > 0)
                       gs0=gs0*exp(cmplx(0,v%z%d*sqrt(sqa0)))
                       gr0=gr0*exp(cmplx(0,-v%z%d*sqrt(sqa0)))
                   elsewhere
                       gs0=gs0*exp(-sqrt(-sqa0)*v%z%d)
                       gr0=gr0*exp(-sqrt(-sqa0)*v%z%d)
                   end where evanescent1


! Operations in x            

                   tfield=gs0; call sfftw_execute(planI); ws=tfield/sqrt(1.*kxm%n)
                   tfield=gr0; call sfftw_execute(planI); wr=tfield/sqrt(1.*kxm%n)

            else;                                                                                       ! v(z) condition
                velref: do iref=1,nrefer

                    gs(:,iref)=ws;  gr(:,iref)=wr; 
!                    do ikx = 1,kxm%n;                                                                         ! wavenumber loop
!      phase shift
!                        sqa = rec%w%all(iwa)**2 / vref(iref,iz,ixs1)**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                                                                                    ! end wavenumber loop

                      sqa0 = rec%w%all(iwa)**2 / vref(iref,iz,ixs1)**2 - kxm%all**2;
                      evanescent2:where (sqa0 > 0)
                          gs(:,iref)=gs(:,iref)*exp(cmplx(0,v%z%d*sqrt(sqa0)))
                          gr(:,iref)=gr(:,iref)*exp(cmplx(0, -v%z%d*sqrt(sqa0)))
                      elsewhere
                          gs(:,iref)=gs(:,iref)*exp(-sqrt(-sqa0)*v%z%d)
                          gr(:,iref)=gr(:,iref)*exp(-sqrt(-sqa0)*v%z%d)
                      end where evanescent2

! Operations in x            

                      tfield=gs(:,iref); call sfftw_execute(planI); gs(:,iref)=tfield/sqrt(1.*kxm%n)
                      tfield=gr(:,iref); call sfftw_execute(planI); gr(:,iref)=tfield/sqrt(1.*kxm%n)


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

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

!                do ix=1,kxm%n
!                    do iref=1,nrefer-1
!                        if (vref(iref,iz,ixs1) .ne. vref(iref+1,iz,ixs1)) then
!                            if (vel(iz,ix) >= vref(iref,iz,ixs1) .and. vel(iz,ix) < vref(iref+1,iz,ixs1)) then
!                                f=(vel(iz,ix) - vref(iref,iz,ixs1))/vel(iz,ix)
!                                g=1-f
!                                ws(ix)=g*gs(ix,iref)+f*gs(ix,iref+1)
!                                wr(ix)=g*gr(ix,iref)+f*gr(ix,iref+1)
!!                                ws(ix)=gs(ix,iref)
!!                                wr(ix)=gr(ix,iref)
!                            elseif (vel(iz,ix) >= vref(nrefer,iz,ixs1)) then
!                                ws(ix)=gs(ix,nrefer)
!                                wr(ix)=gr(ix,nrefer)
!                            end if
!                        else
!                            ws(ix)=gs(ix,iref)
!                            wr(ix)=gr(ix,iref)
!                        end if                            
!                    end do
!                end do
           end if vxz
!        if (mod(iwa,10)==1) write(0,*)  "1 ",iwa,maxval(cabs(ws)),maxval(cabs(wr))
!
! 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 .or. ixw-ihso > kxm%n .or. ixw+ihso < 1 .or. ixw+ihso > kxm%n) then
                    else
                         GG(iz,ish,ixw)=GG(iz,ish,ixw)+conjg(ws(ixw-ihso))*(wr(ixw+ihso))             ! correlation to image
                    end if
                end do
            end do

        end do depth


!write(0,*) "PSPI - passei depth"
end subroutine

end module
