module OW_wei

use sep
use OWvelo
use OW_type
use OW_image
use OW_fftw
use OW_grabvel
use OW_parms

integer,private					:: isp,ni,iwa,nwb,node,inode
real,private					:: wb

contains

!---------------------------------------------------------
subroutine owwei_init(ni_in,isp_in,iwa_in,wb_in,nwb_in,node_in,inode_in)

integer						:: isp_in,ni_in,iwa_in,nwb_in,node_in,inode_in
real						:: wb_in

isp=isp_in
ni=ni_in
iwa=iwa_in
wb=wb_in
nwb=nwb_in
node=node_in
inode=inode_in

end subroutine
!---------------------------------------------------------
subroutine owwei(taper,recv,source,vsal,vel,GG,rwr)

optional					:: rwr
integer						:: iz,ikx,iref,ixw
real						:: sqa,f,g
real						:: minv, maxv
real						:: taper(:),vel(:,:),vsal
real, allocatable				:: vref(:),sqa0(:)
complex, allocatable				:: ws(:),wr(:),gs0(:),gr0(:)
complex, allocatable				:: gs(:,:),gr(:,:),tfld(:)
complex						:: recv(:,:),source(:,:,:),GG(:,:,:),rwr(:,:)

if (allocated(wr)) deallocate(wr);allocate(wr(rec%h%n));		wr=0
if (allocated(ws)) deallocate(ws);allocate(ws(rec%h%n));		ws=0
if (allocated(gs)) deallocate(gs);allocate(gs(rec%h%n,nref));		gs=0
if (allocated(gr)) deallocate(gr);allocate(gr(rec%h%n,nref));		gr=0
if (allocated(gs0)) deallocate(gs0);allocate(gs0(rec%h%n));		gs0=0
if (allocated(gr0)) deallocate(gr0);allocate(gr0(rec%h%n));		gr0=0
if (allocated(sqa0)) deallocate(sqa0);allocate(sqa0(rec%h%n));		sqa0=0
if (allocated(vref)) deallocate(vref);allocate(vref(nref));		vref=0
if (allocated(tfld)) deallocate(tfld);allocate(tfld(rec%h%n));		tfld=0

!!--------------------------------------------------
!! Setting FFTW plans
!!

call owfftw(tfld)

ws(isp)=source(1,iwa,1) ! to be modified to handle generalized sources
wr=recv(:,iwa) 

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

!--------------------------------------------------
! Compute vref
!
            call owvref(nref,vsal,vel(iz,:),vref,minv,maxv)

!!!!!!!
            tfld=ws; call sfftw_execute(planF); ws=tfld/sqrt(1.*kxm%n)
            tfld=wr; call sfftw_execute(planF); wr=tfld/sqrt(1.*kxm%n)
write(0,*) "SUM",sum(tfld),sum(wr)

           vxz:if (minv .eq. maxv) then
                gs0=ws;  gr0=wr;

!      Operations in the wavenumber domain
!      phase shift

                wavenumber1:do ikx = 1,kxm%n

                   sqa = rec%w%all(iwa)**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 sfftw_execute(planI); ws=tfld/sqrt(1.*kxm%n)
                tfld=gr0; call sfftw_execute(planI); wr=tfld/sqrt(1.*kxm%n)

            else
                velref: do iref=1,nref                 
                    gs(:,iref)=ws;  gr(:,iref)=wr; 
                    wavenumber2: do ikx = 1,kxm%n
!      phase shift
                        sqa = rec%w%all(iwa)**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 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(iwa))*(1./vel(iz,:)-1./vref(iref)));
                      gr(:,iref) = gr(:,iref) * exp(cmplx(0,-v%z%d*rec%w%all(iwa))*(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
                    do iref=1,nref-1

                        if (vref(iref) .ne. vref(iref+1)) then
                            if (vel(iz,ikx) >= vref(iref) .and. vel(iz,ikx) < vref(iref+1)) then
                                f=(vel(iz,ikx) - vref(iref))/vel(iz,ikx)
                                g=1-f
                                ws(ikx)=g*gs(ikx,iref)+f*gs(ikx,iref+1)
                                wr(ikx)=g*gr(ikx,iref)+f*gr(ikx,iref+1)
                            end if
                            if (vel(iz,ikx) >= vsal) then ! ref(nref)) then
                                ws(ikx)=gs(ikx,nref)
                                wr(ikx)=gr(ikx,nref)
                            end if
                        else
                            ws(ikx)=gs(ikx,iref)
                            wr(ikx)=gr(ikx,iref)
                        end if                            
                    end do
                end do
           end if vxz

           ws=ws*taper
           wr=wr*taper

!           if (maxval(cabs(wr)).ne.0. .or. maxval(cabs(ws)).ne.0.) write(0,*) iwa,iz, maxval(cabs(wr)),maxval(cabs(ws))

           call owimage(ws,wr,GG(iz,:,:))
           if(exist_file("wr")) then
               rwr(iz,:)=ws
           end if
        end do depth

end subroutine

end module
