module OW_ctrillwei

use sep
use OWvelo_areal
use OW_type
use OW_fftw_areal
use OW_parms_areal
use OW_exec_time
use OW_combine
use wei_myfft_mod
!use OW_pickbp

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,wavelet,vref1,igap)

integer,optional					:: igap
real,optional					:: vref1(:,:,:),rand1(:,:,:),vsal
integer						:: nhorz,n,wavz
character*80					:: hor1,hor2
integer						:: iw,iz,ikx,iref,ixw,inode,nrefvel,i,ix
integer						:: indref(1),sign,ih,jxs,ixs,iref1,ikx1,ihor
real						:: minv,maxv,sqa,rnd_phase,start,finish,sqa1
real						:: vel(:,:),taper(:),f,g,kz2_max,kz2,kz
real, allocatable					:: vref(:),a(:),kz0(:,:)
integer,allocatable					:: indw(:,:)
complex, allocatable				:: ws(:,:),gs0(:,:),gs(:,:),inset(:,:,:)
real, allocatable					:: intmp(:,:,:),temp(:,:,:),pick1(:,:),pick2(:,:)
complex, allocatable				:: tflds(:,:),refl(:,:)
complex						:: GG(:,:),wavelet(:)
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
allocate(kz0(kxm%n,nodes));kz0=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
   allocate(temp(sog%z%n,sog%sh%n,sog%xm%n));intmp=0.
   call sseek("sou_cig",0,0)
   call sreed("sou_cig",intmp,4*sog%z%n*sog%sh%n*sog%xm%n)
   do ix=1,sog%xm%n
      do ih=1,sog%sh%n
         temp(:,ih,ix)=intmp(:,sog%sh%n-ih+1,ix)
      end do
   end do
   intmp=temp
   deallocate(temp)
end if
!--------------------------------------------------
! Setting FFTW plans
allocate(tflds(kxm%n,nodes));tflds=0
call owfftws(tflds,nodes)

if (encode) then
   allocate(inset(nz,nx,sou%w%n));inset=0
else
   allocate(inset(nz,nx,2));inset=0
end if

if (present(rand1).and.present(igap)) then
   call owcombine(intmp,inset,nodes,rand1,igap)
else if (present(rand1).and.not(present(igap))) then
   call owcombine(intmp,inset,nodes,rand1)
else if (not(encode)) then
   call owcombine(intmp=intmp,inset=inset,igap=igap)
end if

deallocate(intmp)

kz2_max =  cos(maxang*4.*atan(1.)/180.)**2
!$OMP parallel do private (inode,&
                              iw,&
                              iz,&
                             ikx,&
                            ikx1,&
                            iref,&
                           iref1,&
                             sqa,&
                            sqa1,&
                            minv,&
                            maxv,&
                              gs,&
                            vref,&
                               f,&
                               g,&
                             kz2,&
                              kz,&
                          indref,&
                          nrefvel) 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=nzmax,nzmin,-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/vref1(:,1,iz)
               if (nrefvel.eq.1) minv=vref(1)
            else
               nrefvel=nref
               if (maxval(vel).ge.vsal) then
                  if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel+1))
                  if (allocated( gs )) deallocate( gs ); allocate(gs(kxm%n,nrefvel+1)); gs=0.
                  vref=0
                  call owvref(nrefvel,vel(iz,:),vref(1:nrefvel),minv,maxv)
                  vref(nref+1)=maxval(vel(iz,:))
                  if (minv==maxv) nrefvel=1
               else
                  if (allocated(vref)) deallocate(vref); allocate(vref(nrefvel))
                  if (allocated( gs )) deallocate( gs ); allocate(gs(kxm%n,nrefvel)); gs=0.
                  vref=0
                  call owvref(nrefvel,vel(iz,:),vref,minv,maxv)
                  if (minv==maxv) nrefvel=1
               end if
            end if

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

              tflds(:,inode)=0.  
              if (encode) then
                 tflds(:,inode)=inset(iz-nzmin+1,:,iw)
              else
                 tflds(:,inode)=inset(iz-nzmin+1,:,1)
              end if
              call owfftexes(tflds(:,inode),inode)
              refl(:,inode)=tflds(:,inode)/sqrt(1.*kxm%n)!*cabs(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)

!! Amplitude correction
!              if (iz==1) then
!                 kz0(:,inode)=sou%w%all(iw)**2 / minv**2 - kxm%all**2
!                 where (kz0(:,inode) > 0 )
!                    ws(:,inode)=ws(:,inode)/kz0(:,inode)
!                 elsewhere
!                    ws(:,inode)=0
!                 end where
!              end if

           if (nrefvel.eq.1) 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
                 kz2 = sqa*minv**2/sou%w%all(iw)**2
                 if (kz2 >= 0. .and. kz2 < kz2_max) then
                    gs0(ikx,inode)=0.
                 else if (kz2 >= kz2_max) then
                    sqa1=sqrt(sqa)*minv/sou%w%all(iw)
                    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
                    sqa = sou%w%all(iw)**2 / vref(iref)**2 - kxm%all(ikx)**2;
                    if (sqa > 0) then
                       sqa1=sqrt(sqa)*vref(iref)/sou%w%all(iw)
                       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 
                 tflds(:,inode)=gs(:,iref)
                 call owifftexes(tflds(:,inode),inode)
                 gs(:,iref)=tflds(:,inode)/sqrt(1.*kxm%n)
                 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 iref1=2,nrefvel
                    where (vel(iz,:) >= 1/vref1(iref1,2,iz) .and. vel(iz,:) < 1/vref1(iref1-1,2,iz))
                      ws(:,inode)=gs(:,iref1)
                    end where
                 end do
                 where (vel(iz,:) > 1/vref1(1,2,iz))
                    ws(:,inode)=gs(:,1)
                 end where
              else
!                 do iref1=1,nrefvel-1
!                    where (vel(iz,:) >= 1/vref1(iref1,2,iz) .and. vel(iz,:) < 1/vref1(iref1+1,2,iz))
!                      ws(:,inode)=gs(:,iref1)
!                    end where
!                 end do
!                 where (vel(iz,:) > 1/vref1(nrefvel,2,iz))
!                    ws(:,inode)=gs(:,nrefvel)
!                 end where
                 do ikx1=1,kxm%n
                    indref=0
                    indref=maxloc(vref, mask = int(vref).lt.int(vel(iz,ikx1))) 
                    if (indref(1).ne.0.and.indref(1).ne.size(vref)-1) then
                       f=(vel(iz,ikx1) - vref(indref(1)))/vel(iz,ikx1)
                       g=1-f
                       ws(ikx1,inode)=g*gs(ikx1,indref(1))+f*gs(ikx1,indref(1)+1)
                    else
                       ws(ikx1,inode)=gs(ikx1,indref(1)+1)
                    end if
                 end do
              end if
           end if

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

!           if ((iz==nzmax-2.or.iz==nzmin).and.inode==2) then
!              write(0,*) iz,iw,ws(:,inode) !minval(abs(ws(:,inode))),maxval(abs(ws(:,inode)))
!           end if

      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)
!    write(5,*) iw,inode,GG(iw,:) !minval(abs(ws(:,inode))),maxval(abs(ws(:,inode)))

! !$OMP END ORDERED
end do

!$OMP END PARALLEL DO
!write(0,*) igap,sign,ws
call owfftdestroy()
deallocate(tflds)
deallocate(ws,gs0,refl,inset)
if (allocated(vref)) deallocate(vref)
end subroutine

end module
