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,ihor,ireal

! use OW_type
! use OW_parms_areal
! use OW_exec_time
! use OW_taper_areal
! use OW_ctrillwei
! use OW_rndseed1
! use OW_encoding
! use OW_shaper
! use lloyd_vel
!
! implicit none
!
! real						:: start,finish,start1,finish1,a(2),b,c
! integer						:: i,nv,ixs,j,ni,nsy,iw,ih,icomb,ix,stat,iz,ihor,n_enc_sog,iq,nq
! integer						:: ind_trac,i_pos,i_ini,i_fin,igap
! integer						:: jgap,jxs,nnn,hij,k,ib,ic,ireal,igap1
! real, allocatable					:: velint(:,:,:),taper(:,:),vel(:,:,:)
! real,allocatable					:: randz(:,:,:),vref(:,:,:)
! complex,allocatable				:: source(:),GG(:,:,:)
! complex, allocatable				:: GGG(:,:,:,:),tfld(:),wave_R(:,:,:),wave_S(:,:,:)
! complex						:: phase
!
!!-----------------------------------------
!! Starting ...
!!




contains

!--------------------------------------------------------------------
subroutine owctrlwei_init(ihor_in,ireal_in)

integer                                                     :: ihor_in,ireal_in
integer, external 					:: omp_get_thread_num,omp_get_num_threads

ihor=ihor_in
ireal=ireal_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()

integer						:: igap
real,optional					:: vref1(:,:,:),rand1(:,:),vsal
integer						:: 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
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(:,:),bws(:,:,:)
complex						:: GG(:,:),wavelet(:)
integer, external 					:: omp_get_thread_num,omp_get_num_threads

!--------------------------------------------------
! Computing taper
!
 allocate(taperx(kxm%n))
 call owtaper(taperx,kxm%n)
 allocate(tapery(kxm%n))
 call owtaper(tapery,kym%n)
!--------------------------------------------------
! Read source
!
 allocate(source(sou%w%n))
 call sreed("sou",source,8*sou%w%n)
!--------------------------------------------------
call ow_blk()
if (encode.and.rand) then
   do ireal=1,nreal
      do igap=1,ngap
         wave_S=0;wave_R=0
         do ihor=1,nhorz
            do izb=1,zblk%nb
               if(zblk%nb_rem /= 0.and.izb==zblk%nb) then
                  zn=zblk%nb_rem
               else
                  zn=zbn
               end if   
!--------------------------------------------------
!
! Read rotated CIGs
!
               allocate(intmp(sog%xm%n,sog%ym%n,sog%shx%n,sog%shy%n,zn));intmp=0.
               call ow_cig_io(zn,ihor,gat,intmp) 
               do iwb=1,wblk%nb


############################################################
############################################################
############################################################

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 velocities
!
 allocate(velint(v%xm%n,v%ym%n,v%z%n))
 call sreed("vel",velint,v%z%n*v%xm%n*v%ym%n*4)

 nv=int((xmin-v%xm%o)/v%xm%d)+1.5
 allocate(vel(sog%z%n,nx))
 vel(1:sog%z%n,:)=velint(1:sog%z%n,nv:nv+nx-1)
 deallocate(velint)
 allocate(velint(v%z%n*v%xm%n),vel(img%z%n*kxm%n),vel0(img%z%n*img%xm%n))
 if (lloyd) then 
   allocate(vref(nref,2,img%z%n))
   call lloyd_vel1(1/vel,vref)
 end if
!--------------------------------------------------
! Wavefield extrapolation
!
 allocate(GG(sou%w%n,kxm%n))
 if (encode.and.rand) then
    allocate(wave_S(sou%w%n,kxm%n),wave_R(sou%w%n,kxm%n))
    do ireal=1,nreal
       do igap=1,ngap
          wave_S=0;wave_R=0
          do ihor=1,nhorz
             call init_random_number(randz)
             if (lloyd) then
             call owexec_time(start)
             GG=0
             call owctrlwei_init(ihor,ireal)
             call owctrlwei(taper=taper,vsal=vsal,vel=vel,GG=GG,sign=+1,&
                            rand1=randz(:,:,ihor,(ireal-1)*ngap+igap),wavelet=source,vref1=vref,igap=igap)
             wave_R=wave_R+GG
             call owexec_time(finish)
             write(0,4) "Receiver wavefield   - realization # ",ireal,' FROM ',nreal,&
                        ' igap ',igap,' FROM ',ngap,' horz ',ihor,' FROM ',nhorz,' in ',finish-start,' s'
             call owexec_time(start)
             GG=0
             call owctrlwei_init(ihor,ireal)
             call owctrlwei(taper=taper,vsal=vsal,vel=vel,GG=GG,sign=-1,&
                            rand1=randz(:,:,ihor,(ireal-1)*ngap+igap),wavelet=source,vref1=vref,igap=igap)
             wave_S=wave_S+GG
             call owexec_time(finish)
             write(0,4) "Source wavefield   - realization # ",ireal,' FROM ',nreal,&
                        ' igap ',igap,' FROM ',ngap,' horz ',ihor,' FROM ',nhorz,' in ',finish-start,' s'
          else
             call owexec_time(start)
             GG=0
             call owctrlwei_init(ihor,ireal)
             call owctrlwei(taper=taper,vsal=vsal,vel=vel,GG=GG,sign=+1,&
                            rand1=randz(:,:,ihor,(ireal-1)*ngap+igap),wavelet=source,igap=igap)
             wave_R=wave_R+GG
             call owexec_time(finish)
             write(0,4) "Receiver wavefield   - realization # ",ireal,' FROM ',nreal,&
                        ' igap ',igap,' FROM ',ngap,' horz ',ihor,' FROM ',nhorz,' in ',finish-start,' s'
             call owexec_time(start)
             GG=0
             call owctrlwei_init(ihor,ireal)
             call owctrlwei(taper=taper,vsal=vsal,vel=vel,GG=GG,sign=-1,&
                            rand1=randz(:,:,ihor,(ireal-1)*ngap+igap),wavelet=source,igap=igap)
             wave_S=wave_S+GG
             call owexec_time(finish)
             write(0,4) "Source wavefield   - realization # ",ireal,' FROM ',nreal,&
                        ' igap ',igap,' FROM ',ngap,' horz ',ihor,' FROM ',nhorz,' in ',finish-start,' s'
4  format(A35,i4,5(A6,i4),A4,f6.2,a2)
          end if
          end do
          call owshaper(source,wave_S,-1)
          call owshaper(source,wave_R,+1)
          call srite("rec_out",wave_R(:,padtraces+tapertr+1:kxm%n-(padtraces+tapertr)),8*sou%w%n*(kxm%n-2*(padtraces+tapertr)))
          call srite("sou_out",wave_S(:,padtraces+tapertr+1:kxm%n-(padtraces+tapertr)),8*sou%w%n*(kxm%n-2*(padtraces+tapertr)))
       end do
    end do
else if (encode.and.gold.or.encode.and.sgnpi) then
    allocate(wave_S(sou%w%n,kxm%n),wave_R(sou%w%n,kxm%n))
    do igap=1,ngap
      wave_S=0;wave_R=0
      do ihor=1,nhorz
      if (lloyd) then
         GG=0
         call owexec_time(start)
         call owctrlwei_init(ihor,1)
         call owctrlwei(taper=taper,vsal=vsal,vel=vel,GG=GG,sign=+1,rand1=randz(:,:,ihor,1),wavelet=source,vref1=vref,igap=igap)
         call owexec_time(finish)
         wave_R=wave_R+GG
             write(0,4) "Receiver wavefield   - realization # ",ireal,' FROM ',nreal,&
                        ' igap ',igap,' FROM ',ngap,' horz ',ihor,' FROM ',nhorz,' in ',finish-start,' s'
         GG=0
         call owexec_time(start)
         call owctrlwei_init(ihor,1)
         call owctrlwei(taper=taper,vsal=vsal,vel=vel,GG=GG,sign=-1,rand1=randz(:,:,ihor,1),wavelet=source,vref1=vref,igap=igap)
         call owexec_time(finish)
             write(0,4) "Source wavefield   - realization # ",ireal,' FROM ',nreal,&
                        ' igap ',igap,' FROM ',ngap,' horz ',ihor,' FROM ',nhorz,' in ',finish-start,' s'
         wave_S=wave_S+GG
      else
         GG=0
         call owexec_time(start)
         call owctrlwei_init(ihor,1)
         call owctrlwei(taper=taper,vsal=vsal,vel=vel,GG=GG,sign=+1,rand1=randz(:,:,ihor,1),wavelet=source,igap=igap)
         call owexec_time(finish)
         wave_R=wave_R+GG
             write(0,4) "Receiver wavefield   - realization # ",ireal,' FROM ',nreal,&
                        ' igap ',igap,' FROM ',ngap,' horz ',ihor,' FROM ',nhorz,' in ',finish-start,' s'
         GG=0
         call owexec_time(start)
         call owctrlwei_init(ihor,1)
         call owctrlwei(taper=taper,vsal=vsal,vel=vel,GG=GG,sign=-1,rand1=randz(:,:,ihor,1),wavelet=source,igap=igap)
         call owexec_time(finish)
         wave_S=wave_S+GG
             write(0,4) "Source wavefield   - realization # ",ireal,' FROM ',nreal,&
                        ' igap ',igap,' FROM ',ngap,' horz ',ihor,' FROM ',nhorz,' in ',finish-start,' s'
      end if
      end do
      call owshaper(source,wave_S,-1)
      call owshaper(source,wave_R,+1)
      call srite("rec_out",wave_R(:,padtraces+tapertr+1:kxm%n-(padtraces+tapertr)),8*sou%w%n*(kxm%n-2*(padtraces+tapertr)))
      call srite("sou_out",wave_S(:,padtraces+tapertr+1:kxm%n-(padtraces+tapertr)),8*sou%w%n*(kxm%n-2*(padtraces+tapertr)))
   end do
!else if (.not.encode) then
!   do igap=1,ngap
!      GGS=0;GGR=0
!      if (lloyd) then
!         call owexec_time(start)
!         call owctrlwei(taper=taper,vel=vel,GG=GGR,sign=+1,wavelet=source,vref1=vref,igap=igap)
!         call owexec_time(finish)
!         write(0,3) "Receiver wavefield - super  data # ",igap,' FROM ',ngap,' in ',finish-start,' s'
!         call owexec_time(start)
!         call owctrlwei(taper=taper,vel=vel,GG=GGS,sign=-1,wavelet=source,vref1=vref,igap=igap)
!         call owexec_time(finish)
!         write(0,3) "Source wavefield   - super  data # ",igap,' FROM ',ngap,' in ',finish-start,' s'
!      else
!         call owexec_time(start)
!         call owctrlwei(taper=taper,vsal=vsal,vel=vel,GG=GGR,sign=+1,wavelet=source,igap=igap)
!         call owexec_time(finish)
!         write(0,3) "Receiver wavefield - super  data # ",igap,' FROM ',ngap,' in ',finish-start,' s'
!         call owexec_time(start)
!         call owctrlwei(taper=taper,vsal=vsal,vel=vel,GG=GGS,sign=-1,wavelet=source,igap=igap)
!         call owexec_time(finish)
!         write(0,3) "Source wavefield   - super  data # ",igap,' FROM ',ngap,' in ',finish-start,' s'
!      end if
!
!      call owshaper(source,GGS)
!      call owshaper(source,GGR)
!    
!      call srite("rec_out",GGR,8*sou%w%n*nx)
!      call srite("sou_out",GGS,8*sou%w%n*nx)
!   end do
 end if
3  format(A35,i4,A6,i4,A4,f6.2,a2)

!--------------------------------------------------
! Finishing ...
!
 close(2)
 close(3)
 close(4)
 close(5)

 deallocate(GG,wave_R,wave_S,taper)
 deallocate(vel)

















##########################################################
##########################################################
##########################################################






!--------------------------------------------------
! 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 (encode) then
   call owcombine(intmp_=intmp,inset_=inset,nodes_=nodes,rand1_=rand1,igap_=igap,ireal_=ireal)
else
   call owcombine(intmp_=intmp,inset_=inset,igap_=igap,ireal_=ireal)
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)
              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
                 if (sqa > 0) 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
                       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=abs(vel(iz,ikx1) - vref(indref(1)))/vel(iz,ikx1)
                       f=f-int(f)
                       g=1-f
!                       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 (exist_file("ws").or.exist_file("wr").and.igap==1) then
              bws(:,iz,iw)=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
if (exist_file("ws").or.exist_file("wr").and.igap==1) then
   if (sign.eq.-1) then
      call srite("ws",bws,8*kxm%n*nz*sou%w%n)
   else
      call srite("wr",bws,8*kxm%n*nz*sou%w%n)
   end if
end if
call owfftdestroy()
deallocate(tflds)
deallocate(ws,gs0,refl,inset)
if (allocated(vref)) deallocate(vref)
if (exist_file("ws").or.exist_file("wr")) then
   deallocate(bws)
end if
end subroutine

end module
