module OW_ctrillwei

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, ixs

contains

!--------------------------------------------------------------------
subroutine owctrlwei_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 owctrlwei(taper,vsal,vel,inset,GG,sign,rand1,vref1)

real,optional					:: vref1(:,:,:),rand1(:,:)
integer						:: iw,iz,ikx,iref,ixw,inode,nrefvel,sign
real						:: minv,maxv,vsal,sqa,rnd_phase
real						:: vel(:,:),inset(:,:),taper(:)
real, allocatable					:: vref(:)
integer,allocatable					:: indw(:,:)
complex, allocatable				:: ws(:,:),gs0(:,:),gs(:,:)
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

!--------------------------------------------------
! Setting FFTW plans
allocate(tflds(kxm%n,nodes));tflds=0
call owfftws(tflds,nodes)

!$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)
            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
!-----------------------------------------------------
           if (encode) then
              if (encz) then
                 if (iz < 325) then
                    rnd_phase=rand1(iw,1)
                 else
                    rnd_phase=rand1(iw,2)
                 end if                
              else
                 rnd_phase=rand1(iw,1)
              end if
           end if

           if (encode) then
              if (injection==0) then
                 tflds(:,inode)=0.   
                 tflds(:,inode)=cmplx(inset(iz,:),0)
                 call owfftexes(tflds(:,inode),inode)
                 refl(:,inode)=tflds(:,inode)/sqrt(1.*kxm%n)
                 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)
                 ws(:,inode)=ws(:,inode)*exp(cmplx(0,eps*rnd_phase))
              else
                 tflds(:,inode)=0.   
                 tflds(:,inode)=cmplx(inset(iz,:),0)
                 tflds(:,inode)=tflds(:,inode)*exp(cmplx(0,eps*rnd_phase))
                 call owfftexes(tflds(:,inode),inode)
                 refl(:,inode)=tflds(:,inode)/sqrt(1.*kxm%n)
                 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)
              end if 
           else
              tflds(:,inode)=0.   
              tflds(:,inode)=cmplx(inset(iz,:),0)
              call owfftexes(tflds(:,inode),inode)
              refl(:,inode)=tflds(:,inode)/sqrt(1.*kxm%n)
              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)
           end if

           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)=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)=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," / CIG # ",ixs
5       format(A17,i5,2(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)
if (allocated(vref)) deallocate(vref)
end subroutine

end module
