module OW_ctrillwei1

use sep
use OWvelo_areal
use OW_type
use OW_fftw_areal
use OW_parms_areal
use OW_exec_time
use wei_myfft_mod

implicit none

public					:: owctrlwei_init,owctrlwei
integer,private				:: nodes

contains

!--------------------------------------------------------------------
subroutine owctrlwei_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(0,*) "****************"
write(0,*) "* NUM.THREADS = ",nodes
write(0,*) "****************"

end subroutine

!---------------------------------------------------------
subroutine owctrlwei(taper,vsal,vel,GG,sign,rand1,vref1)

real,optional					:: vref1(:,:,:),rand1(:,:,:)
integer						:: iw,iz,ikx,iref,ixw,inode,nrefvel,sign,ih,jxs,ixs
real						:: minv,maxv,vsal,sqa,rnd_phase,start,finish
real						:: vel(:,:),taper(:)
real, allocatable					:: vref(:)
integer,allocatable					:: indw(:,:)
complex, allocatable				:: ws(:,:),gs0(:,:),gs(:,:),inset(:,:,:)
real, allocatable					:: intmp(:,:,:)
!real,allocatable					:: incig(:,:,:),inrot(:,:,:)
complex, allocatable				:: tflds(:,:),refl(:,:)
complex						:: GG(:,:)
integer, external 					:: omp_get_thread_num,omp_get_num_threads

allocate(ws(kxm%n,nodes)); ws=0.
allocate(gs0(kxm%n,nodes));gs0=0
allocate(refl(kxm%n,nodes));refl=0
!--------------------------------------------------
!
! Read rotated CIGs
!
allocate(intmp(sog%z%n,sog%sh%n,sog%xm%n));intmp=0.

if (sign.eq.1) then
   call sseek("rec_cig",0,0)
   call sreed("rec_cig",intmp,4*sog%z%n*sog%sh%n*sog%xm%n)
else
   call sseek("sou_cig",0,0)
   call sreed("sou_cig",intmp,4*sog%z%n*sog%sh%n*sog%xm%n)
!   allocate(inrot(sog%z%n,sog%sh%n,sog%xm%n));inrot=0.
end if
!do ixs=1,sog%xm%n
!   jxs=ixs+int((sog%xm%o+sog%sh%o-xmin)/sog%xm%d)+0.5
!   if (sign.eq.-1) then
!      do ih=1,sog%sh%n
!         inrot(:,ih,ixs)=intmp(:,sog%sh%n-ih+1,ixs)
!      end do
!      deallocate(intmp)
!   end if
!end do
!--------------------------------------------------
! Setting FFTW plans
allocate(tflds(kxm%n,nodes));tflds=0
call owfftws(tflds,nodes)
allocate(inset(sog%z%n,nx,sou%w%n));inset=0

call owexec_time(start)

!$OMP parallel do private (   iw,&
                              iz,&
                             ixs,&
                            inode) schedule(dynamic,1) num_threads(nodes)
do iw = 1,sou%w%n-subf
   do ixs=1,sog%xm%n
      jxs=ixs+int((sog%xm%o+sog%sh%o-xmin)/sog%xm%d)+0.5
      do iz=sog%z%n,nwbott,-1
!         if (encode.and.encz.and.iz < 400 ) then
!            inset(iz,jxs:jxs+sog%sh%n-1,iw)=inset(iz,jxs:jxs+sog%sh%n-1,iw)+cmplx(intmp(iz,:,ixs),0)*exp(cmplx(0,eps*rand1(iw,ixs,1)))
!         else if (encode.and.encz.and.iz >= 400 ) then
!            inset(iz,jxs:jxs+sog%sh%n-1,iw)=inset(iz,jxs:jxs+sog%sh%n-1,iw)+cmplx(intmp(iz,:,ixs),0)*exp(cmplx(0,eps*rand1(iw,ixs,2)))
!         else if (encode.and.encw) then
            inset(iz,jxs:jxs+sog%sh%n-1,iw)=inset(iz,jxs:jxs+sog%sh%n-1,iw)+cmplx(intmp(iz,:,ixs),0)*exp(cmplx(0,eps*rand1(iw,ixs,1)))
!         else
!            inset(iz,jxs:jxs+sog%sh%n-1,iw)=inset(iz,jxs:jxs+sog%sh%n-1,iw)+cmplx(intmp(iz,:,ixs),0)
         end if
      end do
   end do
end do
!$OMP end parallel do

call owexec_time(finish)

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


do iw = 1,sou%w%n-subf

        inode=omp_get_thread_num()+1; 
        ws(:,inode)=0;
        do iz=sog%z%n,nwbott,-1

!--------------------------------------------------
! Vref
!
            if (present(vref1)) then
               nrefvel=count(mask = vref1(:,1,iz).ne.-1)
               if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel)); vref=0.
               if (allocated( gs )) deallocate( gs ); allocate(gs(kxm%n,nrefvel)); gs=0.
               vref(1:nrefvel)=1/vref1(1:nrefvel,1,iz)
               if (nrefvel.eq.1) minv=minval(vel(iz,:))
            else
               nrefvel=nref
               if (maxval(vel).ge.vsal) then
                  if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel+1))
                  vref=0
                  call owvref(nrefvel,vel(iz,:),vref(1:nrefvel),minv,maxv)
                  vref(nref+1)=maxval(vel)
               else
                  if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
                  vref=0
                  call owvref(nrefvel,vel(iz,:),vref,minv,maxv)
               end if
            end if
!-----------------------------------------------------

              tflds(:,inode)=0.   
              tflds(:,inode)=inset(iz,:,iw)
              call owfftexes(tflds(:,inode),inode)
              refl(:,inode)=tflds(:,inode)/sqrt(1.*kxm%n)*wavelet(iw)
              tflds(:,inode)=0.
              tflds(:,inode)=ws(:,inode)
              call owfftexes(tflds(:,inode),inode)
              ws(:,inode)=tflds(:,inode)/sqrt(1.*kxm%n)
              ws(:,inode)=ws(:,inode)+refl(:,inode)

           if (nrefvel .eq. 1.or.maxv.eq.minv) then
              gs0(:,inode)=ws(:,inode)
!
!      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
                    gs0(ikx,inode)=gs0(ikx,inode)*exp(cmplx(0,sign*v%z%d*sqrt(sqa)))
                 else
                    gs0(ikx,inode)=0. ! gs0(ikx,inode)*exp(-sqrt(-sqa)*v%z%d)
                 end if 
              end do 

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

           else
              do iref=1,nrefvel                 
                 gs(:,iref)=ws(:,inode)
                 do ikx = 1,kxm%n
!      phase shift
                    sqa = sou%w%all(iw)**2 / vref(iref)**2 - kxm%all(ikx)**2;
                    if (sqa > 0) then
                       gs(ikx,iref)=gs(ikx,iref)*exp(cmplx(0,sign*v%z%d*sqrt(sqa)));
                    else
                       gs(ikx,iref)=0. ! gs(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)

!      split step
                 gs(:,iref) = gs(:,iref) * exp(cmplx(0,sign*v%z%d*sou%w%all(iw))*(1./vel(iz,:)-1./vref(iref)));

              end do

              if (present(vref1)) then
                 do iref=2,nrefvel
                    where (vel(iz,:) >= 1/vref1(iref,2,iz) .and. vel(iz,:) < 1/vref1(iref-1,2,iz))
                      ws(:,inode)=gs(:,iref)
                    end where
                 end do
                 where (vel(iz,:) > 1/vref1(1,2,iz))
                    ws(:,inode)=gs(:,1)
                 end where
              end if
           end if

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

      end do

        if (mod(iw,50)==1) write(0,5)  "DONE FREQUENCY # ",iw," OUT OF ",sou%w%n
5       format(A17,i5,A8,i5)

! !$OMP ORDERED
    GG(iw,:)=ws(:,inode)
! !$OMP END ORDERED
end do

!$OMP END PARALLEL DO

call owfftdestroy()
deallocate(tflds)
deallocate(ws,gs0,refl,inset)
if (allocated(vref)) deallocate(vref)
end subroutine

end module
