module OW_wei_areal

use sep
use OWvelo_areal
use OW_type
use OW_image_areal
use OW_fftw_areal
use OW_parms_areal

integer,private					:: nodes,ixs

contains

!---------------------------------------------------------
subroutine owwei_init(ixs_in)

integer						:: ixs_in
integer, external 					:: omp_get_thread_num,omp_get_num_threads

ixs=ixs_in
if (node.eq.9999) then
  !$OMP PARALLEL
  nodes = omp_get_num_threads()
  !$OMP END PARALLEL
else
  nodes=node
end if
write(0,*) "****************"
write(0,*) "* NUM.THREADS = ",nodes
write(0,*) "****************"

end subroutine
!---------------------------------------------------------
subroutine owwei(taper,recv,source,vel,image,vref1,GGr,GGs)

complex,optional					:: GGr(:,:),GGs(:,:)
real, optional					:: vref1(:,:,:)
integer						:: iz,ikx,iref,ixw,nrefvel,inode,ish
real						:: sqa,f,g
real						:: minv, maxv
real						:: taper(:),vel(:,:)
real, allocatable					:: vref(:),sqa0(:),GG(:,:,:,:)
complex, allocatable				:: ws(:,:),wr(:,:),gs0(:,:),gr0(:,:)
complex, allocatable				:: gs(:,:),gr(:,:)
complex						:: recv(:,:),source(:,:)
complex, allocatable				:: tflds(:,:),tfldr(:,:)
real,allocatable					:: image(:,:,:)
integer, external 					:: omp_get_thread_num,omp_get_num_threads

allocate(ws(kxm%n,nodes)); ws=0.
allocate(wr(kxm%n,nodes)); wr=0.
allocate(gs0(kxm%n,nodes));gs0=0
allocate(gr0(kxm%n,nodes));gr0=0
!allocate(indw(sou%w%n,2));indw=0
allocate(GG(img%z%n,img%sh%n,kxm%n,nodes));GG=0.
!--------------------------------------------------
! Setting FFTW plans
allocate(tflds(kxm%n,nodes));tflds=0
call owfftws(tflds,nodes)

allocate(tfldr(kxm%n,nodes));tfldr=0
call owfftwr(tfldr,nodes)

!$OMP parallel do private (inode,&
                             iw,&
                             iz,&
                            ikx,&
                           ihso,&
                            ish,&
                            ixw,&
                           iref,&
                            sqa,&
                           minv,&
                           maxv,&
                           vref,&
                             gr,&
                             gs,&
                         nrefvel) ordered schedule(dynamic,1) num_threads(nodes)

do iw = 1,rec%w%n

        inode=omp_get_thread_num()+1; 
        ws(:,inode)=0;wr(:,inode)=0.
        ws(:,inode)=source(:,iw)
        wr(:,inode)=recv(:,iw) 

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

!--------------------------------------------------
! Compute vref
!
            if (present(vref1)) then
               nrefvel=count(mask = vref1(:,1,iz).ne.-1)
               if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
               if (allocated(gs)) deallocate(gs); allocate(gs(kxm%n,nrefvel));
               if (allocated(gr)) deallocate(gr); allocate(gr(kxm%n,nrefvel));
               vref=1/vref1(:,1,iz); gs(:,:)=0; gr(:,:)=0
               if (nrefvel==1) then
                  minv=vref(1);maxv=minv
               else
                  minv=vref(nrefvel);maxv=vref(1)
               end if
            else
               minv=minval(vel(iz,:));maxv=minv;vref=minv
            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)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
!      Propagate wavefields when velocity is constant 
!
          if (minv .eq. maxv) then
               gs0(:,inode)=ws(:,inode)
               gr0(:,inode)=wr(:,inode)
!
!      Operations in the wavenumber domain
!      phase shift
                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

                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
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!      Propagate wavefields with reference velocities
!
                 do iref=1,nrefvel

                    gs(:,iref)=ws(:,inode)
                    gr(:,iref)=wr(:,inode)

                    do ikx = 1,kxm%n
!      phase shift
                        sqa = rec%w%all(iw)**2 / vref(iref)**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

! Operations in x            

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

                    tfldr(:,inode)=gr(:,iref)
                    call owifftexer(tfldr(:,inode),inode)
                    gr(:,iref)=tfldr(:,inode)/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
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
!       Interpolate between wavefields
!
             if (present(vref1)) then
                do iref1=2,nrefvel
                   where (vel(iz,:) >= 1/vref1(iref1,2,iz) .and. vel(iz,:) < 1/vref1(iref1-1,2,iz))
                      ws(:,inode)=gs(:,iref1)
                      wr(:,inode)=gr(:,iref1)
                   end where
                end do
                where (vel(iz,:) > 1/vref1(1,2,iz))
                   ws(:,inode)=gs(:,1)
                   wr(:,inode)=gr(:,1)
                end where
             end if

           end if

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

!           if (exist_file("ws")) then
!              GGr(iz,:)=wr
!              GGs(iz,:)=ws
!           end if

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

!              call owimage(iw,ws(:,inode),wr(:,inode),GG(iz,:,:,inode))
!           end if

        end do depth

        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 ! end omega

!$OMP END PARALLEL DO

call owfftdestroy()
deallocate(tfldr,tflds)

   do i=1,nodes
      !$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(iz,ish,ixw)=image(iz,ish,ixw)+GG(iz,ish,ixw,i)
            end do
         end do
      end do
      !$OMP END PARALLEL DO
   end do

deallocate(GG,ws,wr,gs0,gr0)


end subroutine

!if (nwbott.ne.1) then
!   call owvref(1,minval(vel(nwbott-1,:)),vel(1,:),vref,minv,maxv)
!   tfld=0.; tfld=wr; call sfftw_execute(planF); wr=tfld/sqrt(1.*kxm%n)
!   tfld=0.; tfld=ws; call sfftw_execute(planF); ws=tfld/sqrt(1.*kxm%n)
!!      Operations in the wavenumber domain
!!      phase shift
!   do ikx = 1,kxm%n
!      sqa = rec%w%all(iw)**2 / (minv)**2 - kxm%all(ikx)**2
!      if (sqa > 0) then
!         wr(ikx)=wr(ikx)*exp(cmplx(0,dwbott*sqrt(sqa)))
!         ws(ikx)=ws(ikx)*exp(cmplx(0,-dwbott*sqrt(sqa)))
!      else
!         wr(ikx)=wr(ikx)*exp(-sqrt(-sqa)*dwbott)
!         ws(ikx)=ws(ikx)*exp(-sqrt(-sqa)*dwbott)
!      end if
!   end do
!   tfld=wr; call sfftw_execute(planI); wr=tfld/sqrt(1.*kxm%n)
!   tfld=ws; call sfftw_execute(planI); ws=tfld/sqrt(1.*kxm%n)
!   wr=wr*taper
!   ws=ws*taper
!end if


!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
subroutine owwei1(taper,recv,source,vsal,vel,GGr,GGs)

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(:),GGr(:,:),GGs(:,:)

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

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

call owfftw(tfld)

ws=source
wr=recv 

if (nwbott.ne.1) then
   call owvref(1,minval(vel(nwbott-1,:)),vel(1,:),vref,minv,maxv)
   tfld=0.; tfld=wr; call sfftw_execute(planF); wr=tfld/sqrt(1.*kxm%n)
   tfld=0.; tfld=ws; call sfftw_execute(planF); ws=tfld/sqrt(1.*kxm%n)
!      Operations in the wavenumber domain
!      phase shift
   do ikx = 1,kxm%n
      sqa = rec%w%all(iw)**2 / (minv)**2 - kxm%all(ikx)**2
      if (sqa > 0) then
         wr(ikx)=wr(ikx)*exp(cmplx(0,dwbott*sqrt(sqa)))
         ws(ikx)=ws(ikx)*exp(cmplx(0,-dwbott*sqrt(sqa)))
      else
         wr(ikx)=wr(ikx)*exp(-sqrt(-sqa)*dwbott)
         ws(ikx)=ws(ikx)*exp(-sqrt(-sqa)*dwbott)
      end if
   end do
   tfld=wr; call sfftw_execute(planI); wr=tfld/sqrt(1.*kxm%n)
   tfld=ws; call sfftw_execute(planI); ws=tfld/sqrt(1.*kxm%n)
   wr=wr*taper
   ws=ws*taper
end if

        depth:do iz=nwbott,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)

           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(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 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(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 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 

           end if vxz

           ws=ws*taper
           wr=wr*taper
           GGr(iz,:)=wr
           GGs(iz,:)=ws

        end do depth

end subroutine owwei1

!!---------------------------------------------------------------------------

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
subroutine owwei_interf(taper,recv,source,vsal,vel,GGr,GGs)

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(:),GGr(:,:),GGs(:,:)

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

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

call owfftw(tfld)

ws=source
wr=recv 

if (nwbott.ne.1) then
   call owvref(1,minval(vel(nwbott-1,:)),vel(1,:),vref,minv,maxv)
   tfld=0.; tfld=wr; call sfftw_execute(planF); wr=tfld/sqrt(1.*kxm%n)
   tfld=0.; tfld=ws; call sfftw_execute(planF); ws=tfld/sqrt(1.*kxm%n)
!      Operations in the wavenumber domain
!      phase shift
   do ikx = 1,kxm%n
      sqa = sou%w%all(iw)**2 / (minv)**2 - kxm%all(ikx)**2
      if (sqa > 0) then
         wr(ikx)=wr(ikx)*exp(cmplx(0,dwbott*sqrt(sqa)))
         ws(ikx)=ws(ikx)*exp(cmplx(0,-dwbott*sqrt(sqa)))
      else
         wr(ikx)=wr(ikx)*exp(-sqrt(-sqa)*dwbott)
         ws(ikx)=ws(ikx)*exp(-sqrt(-sqa)*dwbott)
      end if
   end do
   tfld=wr; call sfftw_execute(planI); wr=tfld/sqrt(1.*kxm%n)
   tfld=ws; call sfftw_execute(planI); ws=tfld/sqrt(1.*kxm%n)
   wr=wr*taper
   ws=ws*taper
end if

        depth:do iz=nwbott,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)

           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(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 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(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 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 

           end if vxz

           ws=ws*taper
           wr=wr*taper
           GGr(iz,:)=wr
           GGs(iz,:)=ws

        end do depth

end subroutine owwei_interf
!!--------------------------------------------------






end module
