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

integer,optional					:: igap
real,optional					:: vref1(:,:,:),rand1(:,:,:)
integer						:: nhorz,n,wavz
character*80					:: hor1,hor2
integer						:: iw,iz,ikx,iref,ixw,inode,nrefvel,i
integer						:: indref(1),sign,ih,jxs,ixs,iref1,ikx1,ihor
real						:: minv,maxv,vsal,sqa,rnd_phase,start,finish
real						:: vel(:,:),taper(:),f,g
real, allocatable					:: vref(:),a(:)
integer,allocatable					:: indw(:,:)
complex, allocatable				:: ws(:,:),gs0(:,:),gs(:,:),inset(:,:,:)
real, allocatable					:: intmp(:,:,:),temp(:,:,:),pick1(:,:),pick2(:,:)
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.
!
! Delete when lint is working
!
nhorz=0
call from_param("wavz",wavz,13)
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)

   if (nhorz .ne. 0) then
      call owexec_time(start)
      call from_param('hor1rec',hor1)
      open(51,file=hor1,status='old')
      allocate(pick1(sog%xm%n,sog%sh%n));pick1=9999.
      do n=1,sog%xm%n
!         read(51,*,end=1) pick1(n,39:67)
         read(51,*,end=1) pick1(n,53:69)
      end do
1     close(51)
      call from_param('hor2rec',hor1)
      open(51,file=hor1,status='old')
      allocate(pick2(sog%xm%n,sog%sh%n));pick2=9999.
      do n=1,sog%xm%n
!         read(51,*,end=1) pick2(n,39:67)
         read(51,*,end=2) pick2(n,45:77)
      end do
2     close(51)
      allocate(temp(sog%z%n,sog%sh%n,sog%xm%n));temp=0.
      do ihor=1,nhorz
         do ikx=1,sog%xm%n
            do ih=1,sog%sh%n
               do iz=1,sog%z%n
                  if ((iz .ge. int(pick1(ikx,ih))-wavz .and. iz .le. int(pick1(ikx,ih))+wavz) .or. &
                      (iz .ge. int(pick2(ikx,ih))-wavz .and. iz .le. int(pick2(ikx,ih))+wavz)) then
                      temp(iz,ih,ikx)=intmp(iz,ih,ikx)
                  end if
               end do
            end do
         end do
      end do
      intmp=temp
      call owexec_time(finish)

      if (exist_file("pickr")) then
         call to_history("n1",nhorz,"pickr"); call to_history("d1",1.,"pickr"); call to_history("o1",1.,"pickr")  
         call to_history("n2",sog%xm%n,"pickr"); call to_history("d2",sog%xm%d,"pickr"); call to_history("o2",sog%xm%o,"pickr")
         call to_history("n3",sog%sh%n,"pickr"); call to_history("d3",sog%sh%d,"pickr"); call to_history("o3",sog%sh%o,"pickr")
         call to_history("n4",1,"pickr"); call to_history("d4",1.,"pickr"); call to_history("o4",1.,"pickr")
         call to_history("esize",4,"pickr")
         pick1=pick1*sog%z%d
         call srite("pickr",pick1,4*sog%xm%n*sog%sh%n*nhorz)
      end if

      deallocate(temp,pick1)
      write(0,4)  "FINISHED WINDOWING ",int(finish-start),'(s)'
   end if

   if(exist_file("reccig")) then
     call srite("reccig",intmp,4*sog%z%n*sog%xm%n*sog%sh%n)
   end if

!   allocate(temp(sog%z%n,sog%sh%n,sog%xm%n));temp=0.
!   do ih=1,sog%sh%n
!      temp(:,ih,:)=intmp(:,sog%sh%n-ih+1,:)
!   end do
!   intmp=temp
!   deallocate(temp)

else
   call sseek("sou_cig",0,0)
   call sreed("sou_cig",intmp,4*sog%z%n*sog%sh%n*sog%xm%n)
   
   if (nhorz .ne. 0) then
      call owexec_time(start)
      call from_param('hor1sou',hor1)
      open(51,file=hor1,status='old')
      allocate(pick1(sog%xm%n,sog%sh%n));pick1=9999.
      do n=1,sog%xm%n
!         read(51,*,end=2) pick1(n,39:67)
         read(51,*,end=3) pick1(n,53:69)
      end do
3     close(51)
      call from_param('hor2sou',hor1)
      open(51,file=hor1,status='old')
      allocate(pick2(sog%xm%n,sog%sh%n));pick2=9999.
      do n=1,sog%xm%n
!         read(51,*,end=2) pick2(n,39:67)
         read(51,*,end=4) pick2(n,45:77)
      end do
4     close(51)
      allocate(temp(sog%z%n,sog%sh%n,sog%xm%n));temp=0.
      do ihor=1,nhorz
         do ikx=1,sog%xm%n
            do ih=1,sog%sh%n
               do iz=1,sog%z%n
                  if ((iz .ge. int(pick1(ikx,ih))-wavz .and. iz .le. int(pick1(ikx,ih))+wavz) .or. &
                      (iz .ge. int(pick2(ikx,ih))-wavz .and. iz .le. int(pick2(ikx,ih))+wavz)) then
                     temp(iz,ih,ikx)=intmp(iz,ih,ikx)
                  end if
!                  temp(picks(ihor,ih,ikx)-wavz:picks(ihor,ih,ikx)+wavz,ih,ikx)= &
!                   intmp(picks(ihor,ih,ikx)-wavz:picks(ihor,ih,ikx)+wavz,ih,ikx)
               end do
            end do
         end do
      end do
      intmp=temp
      call owexec_time(finish)

      if (exist_file("picks")) then
         call to_history("n1",nhorz,"picks"); call to_history("d1",1.,"picks"); call to_history("o1",1.,"picks")  
         call to_history("n2",sog%sh%n,"picks"); call to_history("d2",sog%sh%d,"picks"); call to_history("o2",sog%sh%o,"picks")
         call to_history("n3",sog%xm%n,"picks"); call to_history("d3",sog%xm%d,"picks"); call to_history("o3",sog%xm%o,"picks")
         call to_history("n3",1,"picks"); call to_history("d3",1.,"picks"); call to_history("o3",0.,"picks")
         call to_history("esize",4,"picks")
         pick1=pick1*sog%z%d
         call srite("picks",pick1,4*sog%xm%n*sog%sh%n*nhorz)
      end if

      deallocate(temp,pick1)
      write(0,4)  "FINISHED WINDOWING ",int(finish-start),'(s)'
4 format(A19,i5,A4)
   end if

   if(exist_file("soucig")) then
     call srite("soucig",intmp,4*sog%z%n*sog%xm%n*sog%sh%n)
   end if
   
!   allocate(temp(sog%z%n,sog%sh%n,sog%xm%n));temp=0.
!   do ih=1,sog%sh%n
!      temp(:,ih,:)=intmp(:,sog%sh%n-ih+1,:)
!   end do
!   intmp=temp
!   deallocate(temp)

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

if (encode.and.rand) then
   allocate(inset(nz,nx,sou%w%n));inset=0
   call owexec_time(start)
   !$OMP parallel do private (   iw,&
                                 iz,&
                                ixs,&
                                jxs,&
                               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
!         inset(:,jxs:jxs+sog%sh%n-1,iw)=inset(:,jxs:jxs+sog%sh%n-1,iw)+ &
!                                        cmplx(intmp(nzmin:nzmax,:,ixs),0)*exp(cmplx(0,eps*rand1(iw,ixs,1)))
         do iz=nzmax,nzmin,-1
            inset(iz-nzmin+1,jxs:jxs+sog%sh%n-1,iw)=inset(iz-nzmin+1,jxs:jxs+sog%sh%n-1,iw)+ &
                                                    cmplx(intmp(iz,:,ixs),0)*exp(cmplx(0,eps*rand1(iw,ixs,1)))
         end do
      end do
   end do
   !$OMP end parallel do
   deallocate(intmp)
   call owexec_time(finish)
   write(0,6)  "FINISHED PHASE-ENCODING ",int(finish-start),'(s)'
6  format(A24,i5,A4)
else if (encode.and.gold) then
   allocate(inset(nz,nx,sou%w%n));inset=0
   call owexec_time(start)
   !$OMP parallel do private (   iw,&
                                ixs,&
                                jxs,&
                               inode) schedule(dynamic,1) num_threads(nodes)
   do iw = 1,sou%w%n-subf
      do ixs=igap,sog%xm%n,ngap
         jxs=ixs+int((sog%xm%o+sog%sh%o-xmin)/sog%xm%d)+0.5
         inset(:,jxs:jxs+sog%sh%n-1,iw)=inset(:,jxs:jxs+sog%sh%n-1,iw)+ & 
                                                    cmplx(intmp(nzmin:nzmax,:,ixs),0)*exp(cmplx(0,eps*rand1(iw,ixs,1)))
      end do
   end do
   !$OMP end parallel do
   deallocate(intmp)
   call owexec_time(finish)
   write(0,6)  "FINISHED PHASE-ENCODING ",int(finish-start),'(s)'
else if (.not.encode) then
   allocate(inset(nz,nx,2));inset=0
   do ixs=igap,sog%xm%n,ngap
      jxs=ixs+int((sog%xm%o+sog%sh%o-xmin)/sog%xm%d)+0.5
      inset(:,jxs:jxs+sog%sh%n-1,1)=inset(:,jxs:jxs+sog%sh%n-1,1)+cmplx(intmp(nzmin:nzmax,:,ixs),0)
   end do
   deallocate(intmp)
end if

!$OMP parallel do private (inode,&
                              iw,&
                              iz,&
                             ikx,&
                            ikx1,&
                            iref,&
                           iref1,&
                             sqa,&
                            minv,&
                            maxv,&
                              gs,&
                            vref,&
                               f,&
                               g,&
                          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)
!write(5,*) iw,iz,inode,int(minval(vref)),int(maxval(vref)),nrefvel
               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)*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) 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
                    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=(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

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

end module
