module OW_pspi_areal2

use OW_type
use OW_parms_areal
use OWvelo_areal
use OW_image_areal

implicit none

private						:: ni,ixs,nwb,vsal
integer						:: ni,ixs,nwb
real						:: wb

contains

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

integer						:: ni_in,ixs_in
real						:: wb_in

ni=ni_in
ixs=ixs_in
wb=wb_in
wb=floor(wb/img%z%d)*img%z%d
nwb=floor(wb/img%z%d)+1

end subroutine
!---------------------------------------------------------
subroutine owpspi(taper,recv,source,vsal,vel,G1)

!optional					:: rwr
integer						:: iw,inode,ino,i1,nout,node
integer						:: iz,ikx,iref,ixw
integer*8,allocatable	   			:: planF(:),planI(:)
real						:: sqa,f,g
real						:: minv, maxv
real						:: taper(:),vel(:,:),vsal
complex						:: G1(:,:,:),recv(:,:),source(:,:,:)
complex, allocatable				:: GG(:,:,:,:),rwr(:,:),rws(:,:,:)
real, allocatable				:: vref(:,:),sqa0(:,:)
complex, allocatable				:: ws(:,:),wr(:,:),gs0(:,:),gr0(:,:)
complex, allocatable				:: gs(:,:,:),gr(:,:,:),tfld(:,:)

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

integer, parameter   				:: FFTW_R2HC=0,FTW_HC2R=1,FFTW_DHT=2,FFTW_REDFT00=3,FFTW_REDFT01=4
integer, parameter   				:: FFTW_REDFT10=5,FFTW_REDFT11=6,FFTW_RODFT00=7,FFTW_RODFT01=8,FFTW_RODFT10=9
integer, parameter   				:: FFTW_RODFT11=10,FFTW_FORWARD=-1,FFTW_BACKWARD=+1,FFTW_DESTROY_INPUT=1
integer, parameter   				:: FFTW_UNALIGNED=2,FFTW_CONSERVE_MEMORY=4,FFTW_EXHAUSTIVE=8,FTW_PRESERVE_INPUT=16
integer, parameter   				:: FFTW_PATIENT=32,FFTW_ESTIMATE=64,FFTW_ESTIMATE_PATIENT=128
integer, parameter   				:: FFTW_BELIEVE_PCOST=256,FFTW_DFT_R2HC_ICKY=512,FFTW_NONTHREADED_ICKY=1024
integer, parameter   				:: FFTW_NO_BUFFERING=2048,FFTW_NO_INDIRECT_OP=4096,FFTW_ALLOW_LARGE_GENERIC=8192
integer, parameter   				:: FFTW_NO_RANK_SPLITS=16384,FFTW_NO_VRANK_SPLITS=32768,FFTW_NO_VRECURSE=65536

node=omp_get_num_procs()
allocate(planF(node));	planF=0
allocate(planI(node));	planI=0

if (allocated(GG))	deallocate(GG);		allocate(GG(img%z%n,img%sh%n,img%xm%n,node));	GG=0
if (allocated(wr))	deallocate(wr);		allocate(wr(img%xm%n,node));			wr=0
if (allocated(ws))	deallocate(ws);		allocate(ws(img%xm%n,node));			ws=0
if (allocated(gs))	deallocate(gs);		allocate(gs(img%xm%n,nref,node));		gs=0
if (allocated(gr))	deallocate(gr);		allocate(gr(img%xm%n,nref,node));		gr=0
if (allocated(gs0))	deallocate(gs0);	allocate(gs0(img%xm%n,node));			gs0=0
if (allocated(gr0))	deallocate(gr0);	allocate(gr0(img%xm%n,node));			gr0=0
if (allocated(sqa0))	deallocate(sqa0);	allocate(sqa0(img%xm%n,node));			sqa0=0
if (allocated(vref))	deallocate(vref);	allocate(vref(nref,node));			vref=0
if (allocated(tfld))	deallocate(tfld);	allocate(tfld(img%xm%n,node));			tfld=0

do i1=1,node
	call sfftw_plan_dft(planF(i1),1,kxm%n,tfld(:,i1),tfld(:,i1),fftW_FORWARD,fftW_PATIENT)
	call sfftw_plan_dft(planI(i1),1,kxm%n,tfld(:,i1),tfld(:,i1),fftw_BACKWARD,fftW_PATIENT)
end do

!! $OMP parallel default(shared) 
!! $OMP do private(iw,iz,ikx,minv,maxv,sqa,iref)

!$OMP parallel do default(shared) private(iw,inode,iz,ikx,minv,maxv,sqa,iref)

do iw = 1,rec%w%n

        inode=omp_get_thread_num()+1;

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

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

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

            tfld(:,inode)=ws(:,inode); call sfftw_execute(planF(inode)); ws(:,inode)=tfld(:,inode)/sqrt(1.*kxm%n)
            tfld(:,inode)=wr(:,inode); call sfftw_execute(planF(inode)); wr(:,inode)=tfld(:,inode)/sqrt(1.*kxm%n)

           vxz:if (minv .eq. maxv) 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
                tfld(:,inode)=gs0(:,inode); call sfftw_execute(planI(inode)); ws(:,inode)=tfld(:,inode)/sqrt(1.*kxm%n)
                tfld(:,inode)=gr0(:,inode); call sfftw_execute(planI(inode)); wr(:,inode)=tfld(:,inode)/sqrt(1.*kxm%n)
 
           else
                velref: do iref=1,nref                 
                    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,inode)**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            
                      tfld(:,inode)=gs(:,iref,inode); call sfftw_execute(planI(inode)); gs(:,iref,inode)=tfld(:,inode)/sqrt(1.*kxm%n)
                      tfld(:,inode)=gr(:,iref,inode); call sfftw_execute(planI(inode)); gr(:,iref,inode)=tfld(:,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,inode)));
                      gr(:,iref,inode) = gr(:,iref,inode) * exp(cmplx(0,-v%z%d*rec%w%all(iw))*(1./vel(iz,:)-1./vref(iref,inode)));

                end do velref

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

           end if vxz

           ws(:,inode)=ws(:,inode)*taper
           wr(:,inode)=wr(:,inode)*taper

           call owimage(ws(:,inode),wr(:,inode),GG(iz,:,:,inode))

        end do depth

!$OMP critical
        G1(:,:,:)=G1(:,:,:)+GG(:,:,:,inode)
!$OMP end critical   

        if (mod(iw,10)==1) write(0,5) " FREQUENCY # ",iw," OUT OF ",rec%w%n," / SYNTH.FCT. # ",rec%x%n, &
					" FROM NODE: ",inode

end do
!! $OMP End do
!! $OMP end parallel

!$OMP end parallel do

5 format(A13,i3,A8,i3,A14,i3,A12,i3)

end subroutine owpspi
end module
