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,recset,souset,GG1,GG2,rand1,rand2,vref1)

real,optional					:: rand1(:),rand2(:,:),vref1(:,:,:)
integer						:: iw,indref(1),iz,ikx,iref,ixw,inode,nrefvel
real						:: minv,maxv,vsal,sqa,rnd_phase
real						:: vel(:,:),recset(:,:),souset(:,:),taper(:)
real, allocatable					:: vref(:,:),f(:),g(:)
complex, allocatable				:: ws(:,:),gs0(:,:),gs(:,:,:)
complex, allocatable				:: wr(:,:),gr0(:,:),gr(:,:,:),tfld(:,:)
complex, allocatable				:: tflds(:,:),tfldr(:,:),refl(:,:)
complex						:: GG1(:,:),GG2(:,:)
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(refl(kxm%n,nodes));refl=0
allocate(f(nodes));f=0
allocate(g(nodes));g=0
allocate(vref(nref,nodes)); vref=0.
allocate(gs(kxm%n,nref,nodes)); gs=0.
allocate(gr(kxm%n,nref,nodes)); gr=0.
!allocate(tfld(kxm%n,nodes));tfld=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,&
                           iref,&
                         indref,&
                            ixw,&
                            sqa,&
                           minv,&
                           maxv) ordered schedule(dynamic,1) num_threads(nodes)


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

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

        do iz=sog%z%n,nwbott,-1

!!--------------------------------------------------
!! Compute vref
!!


            if (present(vref1)) then
               nrefvel=count(mask = vref1(:,1,iz).ne.-1)
               vref(:,inode)=0.
               vref(1:nrefvel,inode)=1/vref1(:,1,iz)
               gs(:,:,inode)=0; gr(:,:,inode)=0
               if (nrefvel==1) then
                  minv=vref(1,inode);maxv=minv
               else
                  minv=vref(nrefvel,inode);maxv=vref(1,inode)
               end if
            else
               minv=minval(vel(iz,:));maxv=minv;vref(:,inode)=minv

!               if (maxval(vel(iz,:)).ge.vsal) then
!                  if (1-vsal/maxval(vel(iz,:)).ge.0.1) then 
!                     if (allocated(vref)) deallocate(vref); allocate(vref(nref+2))
!                     if (allocated(gs)) deallocate(gs); allocate(gs(kxm%n,nref+2));
!                     if (allocated(gr)) deallocate(gr); allocate(gr(kxm%n,nref+2));
!                     vref=0; gs(:,:)=0; gr(:,:)=0
!                     call owvref(nref,vel(iz,:),vref(1:nref),minv,maxv)
!                     vref(nref+1)=0.5*(maxval(vel(iz,:))+vsal);vref(nref+2)=maxval(vel(iz,:))
!                  else
!                     if (allocated(vref)) deallocate(vref); allocate(vref(nref+1))
!                     if (allocated(gs)) deallocate(gs); allocate(gs(kxm%n,nref+1));
!                     if (allocated(gr)) deallocate(gr); allocate(gr(kxm%n,nref+1));
!                     vref=0; gs(:,:)=0; gr(:,:)=0
!                     call owvref(nref,vel(iz,:),vref(1:nref),minv,maxv)
!                     vref(nref+1)=maxval(vel(iz,:))
!                  end if
!               else if(maxval(vel(iz,:)).lt.vsal) then
!                  if (allocated(vref)) deallocate(vref); allocate(vref(nref))
!                  if (allocated(gs)) deallocate(gs); allocate(gs(kxm%n,nref));
!                  if (allocated(gr)) deallocate(gr); allocate(gr(kxm%n,nref));
!                  vref=0; gs(:,:)=0; gr(:,:)=0
!                  call owvref(nref,vel(iz,:),vref,minv,maxv)
!               end if
            end if

!write(0,*) "00",iz,iw,vref

            tflds(:,inode)=0.
            tflds(:,inode)=cmplx(souset(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)

            tfldr(:,inode)=0.
            tfldr(:,inode)=cmplx(recset(iz,:),0)
            call owfftexer(tfldr(:,inode),inode)
            refl(:,inode)=tfldr(:,inode)/sqrt(1.*kxm%n)
            tfldr(:,inode)=0.
            tfldr(:,inode)=wr(:,inode)
            call owfftexer(tfldr(:,inode),inode)
            wr(:,inode)=tfldr(:,inode)/sqrt(1.*kxm%n)
            wr(:,inode)=wr(:,inode)+refl(:,inode)


!write(2,*) "passei 00",iw,iz,sum(abs(wr(:,1))),sum(abs(ws(:,1)))

           vxz:if (minv .eq. maxv .and. maxv.eq.vref(1,inode)) then
               gs0(:,inode)=ws(:,inode)
               gr0(:,inode)=wr(:,inode)
!
!      Operations in the wavenumber domain
!      phase shift
                wavenumber1:do ikx = 1,kxm%n
                   sqa = sou%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 wavenumber1

                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)

!write(2,*) "passei 01",iw,iz,sum(abs(wr(:,1))),sum(abs(ws(:,1)))

            else
                velref: do iref=1,nrefvel                 
                    gs(:,iref,inode)=ws(:,inode)
                    gr(:,iref,inode)=wr(:,inode); 
                    wavenumber2: do ikx = 1,kxm%n
!      phase shift
                        sqa = sou%w%all(iw)**2 / vref(iref,inode)**2 - kxm%all(ikx)**2;
                        evanescent2: if (sqa > 0) then
                            gs(ikx,iref,inode)=gs(ikx,iref,inode)*exp(cmplx(0,-v%z%d*sqrt(sqa)));
                            gr(ikx,iref,inode)=gr(ikx,iref,inode)*exp(cmplx(0,v%z%d*sqrt(sqa)));
                        else
                            gs(ikx,iref,inode)=gs(ikx,iref,inode)*exp(-sqrt(-sqa)*v%z%d)
                            gr(ikx,iref,inode)=gr(ikx,iref,inode)*exp(-sqrt(-sqa)*v%z%d)
                        end if evanescent2
                    end do wavenumber2
! Operations in x            

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

                    tfldr(:,inode)=gr(:,iref,inode)
                    call owifftexer(tfldr(:,inode),inode)
                    gr(:,iref,inode)=tfldr(:,inode)/sqrt(1.*kxm%n)

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

!write(2,*) "passei 02",iw,iz,sum(abs(gr(:,1,1))),sum(abs(gs(:,1,1)))

                end do velref
             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,inode)
                     wr(:,inode)=gr(:,iref,inode)
                   end where
                end do
                where (vel(iz,:) > 1/vref1(1,2,iz))
                      ws(:,inode)=gs(:,1,inode)
                      wr(:,inode)=gr(:,1,inode)
                end where
             else

!                do ikx=1,kxm%n
!                    indref=0
!                    indref=maxloc(vref, mask = int(vref).lt.int(vel(iz,ikx))) 
!                    if (indref(1).ne.0.and.indref(1).ne.size(vref)-1) then
!                       f(inode)=(vel(iz,ikx) - vref(indref(1)))/vel(iz,ikx)
!                       g(inode)=1-f(inode)
!                       ws(ikx,inode)=g(inode)*gs(ikx,indref(1))+f(inode)*gs(ikx,indref(1)+1)
!                       wr(ikx,inode)=g(inode)*gr(ikx,indref(1))+f(inode)*gr(ikx,indref(1)+1)
!                    else
!                       ws(ikx,inode)=gs(ikx,indref(1)+1)
!                       wr(ikx,inode)=gr(ikx,indref(1)+1)
!                    end if
!                end do
             end if
           end if vxz

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

!write(2,*) "passei 03",iw,iz,sum(abs(wr(:,1))),sum(abs(ws(:,1)))

           if (encode) then
           if (rand.and.encz.or.gold.and.encz.or.gold.and.encw) then
              rnd_phase=rand2(iz,iw)    
           else
              rnd_phase=rand1(iw)
           end if
           if (gold) then
              wr(:,inode)=wr(:,inode)*exp(cmplx(0,eps*rnd_phase))
              ws(:,inode)=ws(:,inode)*exp(cmplx(0,eps*rnd_phase))
           else
              wr(:,inode)=wr(:,inode)*exp(cmplx(0,eps*pi/2*rnd_phase))
              ws(:,inode)=ws(:,inode)*exp(cmplx(0,eps*pi/2*rnd_phase))
           end if
           end if
      end do

!write(2,*) "passei 04",iw,iz,sum(abs(wr(:,1))),sum(abs(ws(:,1)))

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

!$OMP ORDERED
    GG1(iw,:)=wr(:,inode)
    GG2(iw,:)=ws(:,inode)
!$OMP END ORDERED

end do

!$OMP END PARALLEL DO

call owfftdestroy()
deallocate(tfldr,tflds)
deallocate(ws,gs0,wr,gr0,refl,f,g,vref,gr,gs)
end subroutine

end module
