module OW_wei_areal

use sep
use OW_type
use OW_image_areal
use OW_fftw_areal
use OW_parms_areal
use OW_pspi_areal
use OW_scttr_areal
use OW_cstret

implicit none

integer, private                             :: idip,nodes

contains
!======================================================================
subroutine owwei_init()
integer, external			:: omp_get_thread_num,omp_get_num_threads
if (node.eq.9999) then
  !$OMP PARALLEL
  nodes = omp_get_num_threads()
  !$OMP END PARALLEL
else
  nodes=node
end if
write(3,*) "****************"
write(3,*) "* NUM.THREADS = ",nodes
write(3,*) "****************"

end subroutine
!======================================================================
subroutine owwei(taper,recv,source,vel,image,illum,sref)

real, optional					:: sref(:,:,:)
complex						:: recv(:,:),source(:,:)
integer						:: iwa,iz,inode,i,ixw,ish
real						:: taper(:),vel(kxm%n,img%z%n),a
real						:: image(:,:,:),illum(kxm%n,img%z%n),illumina(kxm%n,img%z%n,nodes)
real						:: GGG(kxm%n,nsecdim,nz,nodes)
complex						:: ws(kxm%n,nodes),wr(kxm%n,nodes)
complex						:: tflds(kxm%n,nodes),tfldr(kxm%n,nodes)
integer, external					:: omp_get_thread_num,omp_get_num_threads

wr=0
ws=0
GGG=0

!--------------------------------------------------
! Setting FFTW plans
tflds=0
call owfftws(tflds,nodes)
tfldr=0
call owfftwr(tfldr,nodes)

!$OMP parallel do private(inode,iwa,iz) ordered schedule(dynamic,1) num_threads(nodes)
   do iwa = 1,rec%w%n
      inode=omp_get_thread_num()+1
      ws(tapertr+padtraces+1:rec%h%n+padtraces+tapertr,inode)=source(:,iwa)
      wr(tapertr+padtraces+1:rec%h%n+padtraces+tapertr,inode)=recv(:,iwa)
      do iz=max(2,nzmin),nzmax
         if (present(sref)) then
            call owpspi( iwa, inode, 1,  ws(:,inode), vel(:,iz-1), tflds(:,inode), "sou", sref(:,:,iz-1))
            call owpspi( iwa, inode,-1,  wr(:,inode), vel(:,iz-1), tfldr(:,inode), "rec", sref(:,:,iz-1))
         else
            call owpspi( iwa, inode, 1,  ws(:,inode), vel(:,iz-1), tflds(:,inode), "sou")
            call owpspi( iwa, inode,-1,  wr(:,inode), vel(:,iz-1), tfldr(:,inode), "rec")
         end if
         ws(:,inode)=ws(:,inode)*taper
         wr(:,inode)=wr(:,inode)*taper
         illumina(:,iz,inode)=illumina(:,iz,inode)+ws(:,inode)*conjg(ws(:,inode))
         call owimage(.false., iwa, ws(:,inode),  wr(:,inode), GGG(:,:,iz-nzmin+1,inode), vel(:,iz-1))
      end do
      if (mod(iwa,50)==1) write(3,6) " FREQUENCY # ",iwa," OUT OF ",rec%w%n," / SYNTH.FCT. # ",rec%xd%n, &
				" FROM NODE: ",inode
   end do
!$OMP END PARALLEL DO
   do i=1,nodes
      !$OMP parallel do private(ixw,iz,ish)
      do iz=1,nz
         illum(:,iz)=illum(:,iz)+illumina(:,iz,i)
         do ish=1,nsecdim
            do ixw=1,kxm%n
               image(ixw,ish,iz)=image(ixw,ish,iz)+GGG(ixw,ish,iz,i)
            end do
         end do
      end do
      !$OMP END PARALLEL DO
   end do
6 format(A13,i3,A8,i3,A14,i3,A12,i3)
call owfftdestroy()
end subroutine
!======================================================================
subroutine owzwht(adj,wgcig)!,ofcig)
logical                                                     :: adj
real                                                        :: wgcig(:,:,:)!,ofcig(:,:,:)
integer                                                     :: iz,ix,ish
real		                                        :: temp(img%z%n)

!if (adj) then
!   temp=0.;wgcig=0.
!   do ish=1,nsecdim;do ix=1,img%xm%n
!      temp=abs(ofcig(ix,ish,:))
!      do iz=5,img%z%n-4
!         temp(iz)=sum(temp(iz-4:iz+4))/9
!      end do
!      wgcig(ix,ish,:)=wgcig(ix,ish,:)+sqrt(real(temp))*ofcig(ix,ish,:)
!   end do;end do
!else
   temp=0.!;ofcig=0.
   do ish=1,nsecdim;do ix=1,img%xm%n
      temp=abs(wgcig(ix,ish,:))
      do iz=5,img%z%n-4
         temp(iz)=sum(temp(iz-4:iz+4))/9
      end do
!      ofcig(ix,ish,:)=ofcig(ix,ish,:)+sqrt(real(temp))*wgcig(ix,ish,:)
      wgcig(ix,ish,:)=temp*wgcig(ix,ish,:)
   end do;end do
!end if
end subroutine
!======================================================================
subroutine owsmooth(adj,wgcig)!,ofcig)
logical                                                     :: adj
real                                                        :: wgcig(:,:,:)!,ofcig(:,:,:)
integer                                                     :: iz,ix,ish
real		                                        :: temp(img%z%n)

!if (adj) then
!   temp=0.;wgcig=0.
!   do ish=1,nsecdim;do ix=1,img%xm%n
!      temp=abs(ofcig(ix,ish,:))
!      do iz=5,img%z%n-4
!         temp(iz)=sum(temp(iz-4:iz+4))/9
!      end do
!      wgcig(ix,ish,:)=wgcig(ix,ish,:)+sqrt(real(temp))*ofcig(ix,ish,:)
!   end do;end do
!else
   temp=0.!;ofcig=0.
   do ish=1,nsecdim;do ix=1,img%xm%n
      temp=abs(wgcig(ix,ish,:))
      do iz=5,img%z%n-4
         temp(iz)=sum(temp(iz-4:iz+4))/9
      end do
!      ofcig(ix,ish,:)=ofcig(ix,ish,:)+sqrt(real(temp))*wgcig(ix,ish,:)
      wgcig(ix,ish,:)=temp
   end do;end do
!end if
end subroutine
!======================================================================
subroutine owxwht(adj,wgcig,ofcig)
logical                                                     :: adj
real                                                        :: wgcig(:,:,:),ofcig(:,:,:)
integer                                                     :: iz,ix,ish
real		                                        :: temp(img%xm%n)

if (adj) then
   temp=0.;wgcig=0.
   do iz=1,img%z%n;do ish=1,nsecdim
      temp=ofcig(:,ish,iz)
      do ix=int(float(nxfilt)/2.)+1,img%xm%n-int(float(nxfilt)/2.)
         temp(ix)=sum(temp(ix-int(float(nxfilt)/2.):ix+int(float(nxfilt)/2.)))/nxfilt
      end do
      wgcig(:,ish,iz)=wgcig(:,ish,iz)+temp
   end do;end do
else
   temp=0.;ofcig=0.
   do iz=1,img%z%n;do ish=1,nsecdim
      temp=wgcig(:,ish,iz)
      do ix=int(float(nxfilt)/2.)+1,img%xm%n-int(float(nxfilt)/2.)
         temp(ix)=sum(temp(ix-int(float(nxfilt)/2.):ix+int(float(nxfilt)/2.)))/nxfilt
      end do
      ofcig(:,ish,iz)=ofcig(:,ish,iz)+temp
   end do;end do
end if
end subroutine
!======================================================================
subroutine owdso(adj,wgcig)!,ofcig)
logical						:: adj
real						:: wgcig(:,:,:)!,ofcig(:,:,:)
integer						:: iz,ix,ish
real						:: h(nsecdim)

   do ish=1,nsecdim
      h(ish)=(ish-1)*dsecdim+osecdim
   end do
   h = (abs(h / maxval(h)))!**2

!if (adj) then
!   wgcig=0.
   do ish=1,nsecdim;do ix=1,img%xm%n
      wgcig(ix,ish,:)=wgcig(ix,ish,:)*h(ish)!*ofcig(ix,ish,:)
   end do;end do
!else
!   ofcig=0.
!   do ish=1,nsecdim;do ix=1,img%xm%n
!      ofcig(ix,ish,:)=ofcig(ix,ish,:)+h(ish)*wgcig(ix,ish,:)
!   end do;end do
!end if
end subroutine

!======================================================================
subroutine owzphase(wgcig)

real						:: wgcig(:,:,:)
complex						:: phase(img%z%n),tfld(img%z%n)
integer						:: iz,ix,ish

phase = exp(cmplx(cos(pi/4.),sin(pi/4.))*sqrt(abs(kz%all)))
write(0,*) phase
call owfftw(tfld)
do ish=1,nsecdim
   do ix=1,img%xm%n
      tfld=cmplx(wgcig(ix,ish,:),0)
      call sfftw_execute(planF)
      tfld=tfld/sqrt(1.*img%z%n)*phase
      call sfftw_execute(planI)
      wgcig(ix,ish,:)=-real(tfld/sqrt(1.*img%z%n))
   end do
end do
call sfftw_destroy_plan(planF)
call sfftw_destroy_plan(planI)
end subroutine
!======================================================================
subroutine owwemva_init(idip_in)
integer      :: idip_in
integer, external			:: omp_get_thread_num,omp_get_num_threads
if (node.eq.9999) then
  !$OMP PARALLEL
  nodes = omp_get_num_threads()
  !$OMP END PARALLEL
else
  nodes=node
end if
write(3,*) "****************"
write(3,*) "* NUM.THREADS = ",nodes
write(3,*) "****************"
idip=idip_in
end subroutine
!======================================================================
subroutine owwemva(adj,taper,recv,source,vel,deltaS,deltaI,sref)

real, optional					:: sref(:,:,:)
logical						:: adj
complex						:: recv(:,:),source(:,:)
integer						:: iwa,iz,inode,i,ixw,ish
real						:: taper(:),vel(kxm%n,img%z%n)
real						:: deltaS(kxm%n,img%z%n)
real						:: deltaI(kxm%n,nsecdim,img%z%n)
real						:: GGG(kxm%n,img%z%n,nodes)
real						:: GG(kxm%n,nsecdim,img%z%n,nodes)
complex						:: ws(kxm%n,nodes),wr(kxm%n,nodes)
complex						:: ddws(kxm%n,nodes),ddwr(kxm%n,nodes)
complex, allocatable				:: dws(:,:,:),dwr(:,:,:)
complex						:: bws(kxm%n,img%z%n,nodes),bwr(kxm%n,img%z%n,nodes)
complex						:: tflds(kxm%n,nodes),tfldr(kxm%n,nodes)
complex						:: dws_z(kxm%n,nodes),dwr_z(kxm%n,nodes)
complex						:: zddws(kxm%n,img%z%n,sou%w%n),zddwr(kxm%n,img%z%n,sou%w%n)
complex						:: ddwsz(kxm%n,img%z%n,sou%w%n),ddwrz(kxm%n,img%z%n,sou%w%n)
integer, external					:: omp_get_thread_num,omp_get_num_threads

wr=0
ws=0
ddwr=0
ddws=0
bwr=0
bws=0

!--------------------------------------------------
! Setting FFTW plans
tflds=0
call owfftws(tflds,nodes)
tfldr=0
call owfftwr(tfldr,nodes)

if (adj) then
   GGG=0
   allocate(dws(kxm%n,img%z%n,nodes),dwr(kxm%n,img%z%n,nodes))
   dwr_z=0
   dws_z=0
   dwr=0
   dws=0
  
   !$OMP parallel do private(inode,iwa,iz) ordered schedule(dynamic,1) num_threads(nodes)
   do iwa = 1,rec%w%n
       inode=omp_get_thread_num()+1
       ws(tapertr+padtraces+1:rec%h%n+padtraces+tapertr,inode)=source(:,iwa)
       wr(tapertr+padtraces+1:rec%h%n+padtraces+tapertr,inode)=recv(:,iwa)
       dws(:,:,inode)=0.
       dwr(:,:,inode)=0.
       bws(:,:,inode)=0.
       bwr(:,:,inode)=0.
       dws_z(:,inode)=0.
       dwr_z(:,inode)=0.
       ddwr(:,inode)=0
       ddws(:,inode)=0
       bws(:,1,inode)=ws(:,inode)
       bwr(:,1,inode)=wr(:,inode)
!!!!!!!!!!!!!!!!!!!!!!!
!
! Background wavefields
! bws and bwr
!

       do iz=nwbott+1,img%z%n
          if (present(sref)) then
             call owpspi(  iwa, inode, 1,  ws(:,inode), vel(:,iz-1), tflds(:,inode), "sou", sref(:,:,iz-1))
             call owpspi(  iwa, inode,-1,  wr(:,inode), vel(:,iz-1), tfldr(:,inode), "rec", sref(:,:,iz-1))
          else
             call owpspi(  iwa, inode, 1,  ws(:,inode), vel(:,iz-1), tflds(:,inode), "sou")
             call owpspi(  iwa, inode,-1,  wr(:,inode), vel(:,iz-1), tfldr(:,inode), "rec")
          end if
          ws(:,inode)=ws(:,inode)*taper
          wr(:,inode)=wr(:,inode)*taper
          bws(:,iz,inode)=ws(:,inode)
          bwr(:,iz,inode)=wr(:,inode)
       end do

!!!!!!!!!!!!!!!!!!!!!!!
!
! Perturbed wavefields
!
       do iz=1,img%z%n
          call owimage(.true.,  iwa, bws(:,iz,inode),  dwr(:,iz,inode), deltaI(:,:,iz), vel(:,iz), .false.)
          call owimage(.true.,  iwa, dws(:,iz,inode),  bwr(:,iz,inode), deltaI(:,:,iz), vel(:,iz), .true. )
       end do
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Total perturbed wavefields
! perturbed + scattered
!
       do iz=img%z%n,1,-1
          dws_z(:,inode)=dws_z(:,inode)+dws(:,iz,inode)
          dws_z(:,inode)=dws_z(:,inode)*taper
          if(present(sref)) then
             call owpspi(  iwa, inode, -1, dws_z(:,inode), vel(:,iz), tflds(:,inode), "sou", sref(:,:,iz))
          else
             call owpspi(  iwa, inode, -1, dws_z(:,inode), vel(:,iz), tflds(:,inode), "sou")
          end if
          ddws(:,inode)=0
          call owscctr(iwa, 1,ddws(:,inode),bws(:,iz,inode),vel(:,iz))
          if (exist_file("ddws").and.idip==int(rec%xd%n/2.)) then
             zddws(:,iz,iwa)=dws_z(:,inode)
             ddwsz(:,iz,iwa)=ddws(:,inode)
          end if
          GGG(:,iz,inode)=GGG(:,iz,inode)+conjg(ddws(:,inode))*dws_z(:,inode)/rec%w%n

          dwr_z(:,inode)=dwr_z(:,inode)+dwr(:,iz,inode)
          dwr_z(:,inode)=dwr_z(:,inode)*taper
          if(present(sref)) then
             call owpspi(  iwa, inode, +1, dwr_z(:,inode), vel(:,iz), tfldr(:,inode), "rec", sref(:,:,iz))
          else
             call owpspi(  iwa, inode, +1, dwr_z(:,inode), vel(:,iz), tfldr(:,inode), "rec")
          end if
          ddwr(:,inode)=0
          call owscctr(iwa,-1,ddwr(:,inode),bwr(:,iz,inode),vel(:,iz))
          if (exist_file("ddwr").and.idip==int(rec%xd%n/2.)) then
             zddwr(:,iz,iwa)=dwr_z(:,inode)
             ddwrz(:,iz,iwa)=ddwr(:,inode)
          end if
          GGG(:,iz,inode)=GGG(:,iz,inode)+conjg(ddwr(:,inode))*dwr_z(:,inode)/rec%w%n
       end do

       if (mod(iwa,50)==1) write(3,6) " FREQUENCY # ",iwa," OUT OF ",rec%w%n," / SYNTH.FCT. # ",rec%xd%n, &
					" FROM NODE: ",inode
   end do
   !$OMP END PARALLEL DO
   if (exist_file("ddws").and.idip==int(rec%xd%n/2.)) then
      call srite("ddws",zddws,8*img%z%n*kxm%n*rec%w%n)
   end if
   if (exist_file("ddwr").and.idip==int(rec%xd%n/2.)) then
      call srite("ddwr",zddwr,8*img%z%n*kxm%n*rec%w%n)
   end if
   if (exist_file("ddwsz").and.idip==int(rec%xd%n/2.)) then
      call srite("ddwsz",ddwsz,8*img%z%n*kxm%n*rec%w%n)
   end if
   if (exist_file("ddwrz").and.idip==int(rec%xd%n/2.)) then
      call srite("ddwrz",ddwrz,8*img%z%n*kxm%n*rec%w%n)
   end if
   call owfftdestroy()
   do i=1,nodes
      !$OMP parallel do private(ixw,iz)
      do iz=1,img%z%n
         do ixw=1,kxm%n
            deltaS(ixw,iz)=deltaS(ixw,iz)+GGG(ixw,iz,i)
         end do
      end do
      !$OMP END PARALLEL DO
   end do
   deallocate(dwr,dws)
else
   GG=0
   allocate(dwr(1,kxm%n,nodes));               dwr=0
   allocate(dws(1,kxm%n,nodes));               dws=0
   !$OMP parallel do private(inode,iwa,iz) ordered schedule(dynamic,1) num_threads(nodes)
    do iwa = 1,rec%w%n
       inode=omp_get_thread_num()+1; 
       ws(tapertr+padtraces+1:rec%h%n+padtraces+tapertr,inode)=source(:,iwa)
       wr(tapertr+padtraces+1:rec%h%n+padtraces+tapertr,inode)=recv(:,iwa)
       dws(1,:,inode)=0
       dwr(1,:,inode)=0
       bws(:,:,inode)=0.
       bwr(:,:,inode)=0.
       bws(:,1,inode)=ws(:,inode)
       bwr(:,1,inode)=wr(:,inode)
!!!!!!!!!!!!!!!!!!!!!!!
!
! Background wavefields
! bws and bwr
!
       do iz=nwbott+1,img%z%n
          if (present(sref)) then
             call owpspi(  iwa, inode, 1,  ws(:,inode), vel(:,iz-1), tflds(:,inode), "sou", sref(:,:,iz-1))
             call owpspi(  iwa, inode,-1,  wr(:,inode), vel(:,iz-1), tfldr(:,inode), "rec", sref(:,:,iz-1))
          else
             call owpspi(  iwa, inode, 1,  ws(:,inode), vel(:,iz-1), tflds(:,inode), "sou")
             call owpspi(  iwa, inode,-1,  wr(:,inode), vel(:,iz-1), tfldr(:,inode), "rec")
          end if
          ws(:,inode)=ws(:,inode)*taper
          wr(:,inode)=wr(:,inode)*taper
          bws(:,iz,inode)=ws(:,inode)
          bwr(:,iz,inode)=wr(:,inode)
       end do
       do iz=1,img%z%n
          ddws(:,inode)=0
          call owscctr(iwa, 1,ddws(:,inode),bws(:,iz,inode),vel(:,iz))
          dws(1,:,inode)=dws(1,:,inode)+ddws(:,inode)*deltaS(:,iz)/rec%w%n
          ddwr(:,inode)=0
          call owscctr(iwa,-1,ddwr(:,inode),bwr(:,iz,inode),vel(:,iz))
          dwr(1,:,inode)=dwr(1,:,inode)+ddwr(:,inode)*deltaS(:,iz)/rec%w%n
          if (present(sref)) then
             call owpspi(  iwa, inode, 1, dws(1,:,inode), vel(:,iz), tflds(:,inode), "sou", sref(:,:,iz))
             call owpspi(  iwa, inode,-1, dwr(1,:,inode), vel(:,iz), tfldr(:,inode), "rec", sref(:,:,iz))
          else
             call owpspi(  iwa, inode, 1, dws(1,:,inode), vel(:,iz), tflds(:,inode), "sou")
             call owpspi(  iwa, inode,-1, dwr(1,:,inode), vel(:,iz), tfldr(:,inode), "rec")
          end if
          dws(1,:,inode)=dws(1,:,inode)*taper
          dwr(1,:,inode)=dwr(1,:,inode)*taper
          call owimage(.false., iwa, dws(1,:,inode), bwr(:,iz,inode), GG(:,:,iz,inode),vel(:,iz))
          call owimage(.false., iwa, bws(:,iz,inode), dwr(1,:,inode), GG(:,:,iz,inode),vel(:,iz))
       end do
       if (mod(iwa,50)==1) write(3,6) " FREQUENCY # ",iwa," OUT OF ",rec%w%n," / SYNTH.FCT. # ",rec%xd%n, &
					" FROM NODE: ",inode
   end do
   !$OMP END PARALLEL DO
   call owfftdestroy()
   do i=1,nodes
      !$OMP parallel do private(ixw,iz,ish)
      do iz=1,img%z%n
         do ish=1,nsecdim
            do ixw=1,kxm%n
               deltaI(ixw,ish,iz)=deltaI(ixw,ish,iz)+GG(ixw,ish,iz,i)
            end do
         end do
      end do
      !$OMP END PARALLEL DO
   end do
   deallocate(dwr,dws)
end if

6 format(A13,i3,A8,i3,A14,i3,A12,i3)
end subroutine
!==================================================
end module
