module OW_migrate

use OW_type
use OW_param

implicit none

contains

subroutine owmigra(nref,wbott,vref,velint,source,image)

integer						:: ixs,iw,iz,ikx,iref,ixw,ish,ino,k,i1,isho
integer						:: nref,nv,ihso,isp,inode,nnode,ni
integer						:: nrec0(3),nrec1(3),nrec2(3),nrec3(3)
real						:: sqa,vref(:,:,:),velint(:,:)
real						:: minv, maxv, dvel, owrk
real, allocatable				:: vel(:,:,:),sqa0(:,:),vwaux(:,:)
complex, allocatable				:: ws(:,:),wr(:,:),gs0(:,:),gr0(:,:),gs(:,:,:),gr(:,:,:)
complex, allocatable				:: recv(:,:,:),GG(:,:,:,:)
complex						:: source(:,:,:),image(:,:,:)
complex, allocatable				:: tfield(:,:),wav(:,:,:)

integer						:: nwbott
real						:: dwbott,vwbott,wbot,wbott

!type (axis) 					:: kxm, h, w, x, xm, z, sh
!type (dat_grid)					:: rec, sou
!type (mod_grid)					:: img
!type (velo)					:: v

integer, external 				:: omp_get_thread_num,omp_get_num_procs,omp_get_num_threads,omp_set_num_threads

integer, parameter   				:: FFTW_R2HC=0,FTW_HC2R=1,FFTW_DHT=2,FFTW_REDFT00=3,FFTW_REDFT01=4
integer, parameter   				:: FFTW_REDFT10=5,FFTW_REDFT11=6,FFTW_RODFT00=7,FFTW_RODFT01=8,FFTW_RODFT10=9
integer, parameter   				:: FFTW_RODFT11=10,FFTW_FORWARD=-1,FFTW_BACKWARD=+1,FFTW_DESTROY_INPUT=1
integer, parameter   				:: FFTW_UNALIGNED=2,FFTW_CONSERVE_MEMORY=4,FFTW_EXHAUSTIVE=8,FTW_PRESERVE_INPUT=16
integer, parameter   				:: FFTW_PATIENT=32,FFTW_ESTIMATE=64,FFTW_ESTIMATE_PATIENT=128
integer, parameter   				:: FFTW_BELIEVE_PCOST=256,FFTW_DFT_R2HC_ICKY=512,FFTW_NONTHREADED_ICKY=1024
integer, parameter   				:: FFTW_NO_BUFFERING=2048,FFTW_NO_INDIRECT_OP=4096,FFTW_ALLOW_LARGE_GENERIC=8192
integer, parameter   				:: FFTW_NO_RANK_SPLITS=16384,FFTW_NO_VRANK_SPLITS=32768,FFTW_NO_VRECURSE=65536
 
!integer*8, pointer,private   			:: planF,planI
!integer*8, private	   				:: planF,planI
integer*8, allocatable 				:: planF(:),planI(:)

 nnode=omp_get_num_threads()
 write(0,*) "01 nref = ", nref

 nnode=2

 write(0,*) "02 nnode = ", nnode

 allocate(GG(img%z%n,img%sh%n,rec%h%n,nnode)); GG=0.

 wbot=floor(wbott/img%z%d)*img%z%d
 nwbott=floor(wbot/img%z%d)+1
 dwbott=wbot/2

nrec0=(/rec%h%n,rec%w%n,rec%x%n/);nrec1=(/rec%h%n,rec%w%n,1/);nrec3=(/1,1,1/)

!--------------------------------------------------
! Setting FFTW plans
!

allocate(tfield(kxm%n,nnode),planF(nnode),planI(nnode))
call sfftw_plan_dft(planF(1),1,kxm%n,tfield(:,1),tfield(:,1),fftW_FORWARD,fftW_PATIENT)
call sfftw_plan_dft(planI(1),1,kxm%n,tfield(:,1),tfield(:,1),fftw_BACKWARD,fftW_PATIENT)
deallocate(tfield)

write(0,*) "FFTplan computed"

do i1=2,nnode
   planF(i1)=planF(1)
   planI(i1)=planI(1)
end do

 if(allocated(recv)) deallocate(recv); allocate(recv(rec%h%n,rec%w%n,nnode))
 if(allocated(vel))  deallocate(vel) ; allocate(vel(v%z%n,rec%h%n,nnode))
 if(allocated(wav))  deallocate(wav) ; allocate(wav(kxm%n,rec%w%n,nnode))
!--------------------------------------------------
! Migrating
!

! OMP parallel firstprivate(wbot,nwbott,dwbott,nnode) num_threads(2) 
write(0,*)"1111111111111111111"
!$OMP parallel default(shared) private(ixs,inode,iw,ikx,iz,iref,ish,isho,ixw,owrk,nrec2,nv,ni,vwbott,sqa) firstprivate(wbot,nwbott,dwbott) num_threads(2)
inode=omp_get_thread_num()+1
!$OMP do
do ixs=1,rec%x%n
write(0,*)"222222222222222222"

write(0,*) "Start do loop shot # ", ixs," from node ",inode

 wav(:,:,inode)=source(:,:,1)
                  
write(0,*) "Start do loop shot # ", ixs," from node ",inode

 nrec2=(/0,0,ixs-1/)
 call sreed_window("in",3,nrec0,nrec1,nrec2,nrec3,8,recv(:,:,inode))
!  call sreed("in",recv(:,:,inode),8*rec%h%n*rec%w%n)

write(0,*) "Read recv from node: ",inode


! workspace origin 
    owrk=(ixs-1)*rec%x%d+rec%x%o+rec%h%o
! grab velocity in the workspace       
    nv = floor((owrk-v%xm%o)/v%xm%d)+1.5
    ni = nv-floor((img%xm%o-v%xm%o)/img%xm%d)+1.5

    vel=0;
    vel(:,:,inode)=velint(:,nv:nv+rec%h%n-1);

!-----------------------------------------------------------
! If required performs water column datuming
!
    if (wbot.ne.0.) then
       if (allocated(vwaux)) deallocate(vwaux);		allocate(vwaux(nwbott,nnode));	vwaux=0.
       do i1=1,nwbott
          vwaux(i1,inode)=sum(vel(i1,:,inode))
       end do
       vwbott=sum(vwaux(:,inode))/(nwbott*kxm%n)
       
       do iw = 1,rec%w%n

        if (allocated(wr)) deallocate(wr);		allocate(wr(rec%h%n,nnode));	wr=0
        if (allocated(gr0)) deallocate(gr0);		allocate(gr0(rec%h%n,nnode));	gr0=0
        if (allocated(ws)) deallocate(ws);		allocate(ws(rec%h%n,nnode));	ws=0
        if (allocated(gs0)) deallocate(gs0);		allocate(gs0(rec%h%n,nnode));	gs0=0
        if (allocated(sqa0)) deallocate(sqa0);		allocate(sqa0(rec%h%n,nnode));	sqa0=0
        if (allocated(tfield)) deallocate(tfield);	allocate(tfield(rec%h%n,nnode));	tfield=0

        ws(:,inode)=wav(:,iw,inode)
        wr(:,inode)=recv(:,iw,inode) 

        do iz=1,3
            tfield=0
            tfield(:,inode)=wr(:,inode); call sfftw_execute(planF(inode)); wr(:,inode)=tfield(:,inode)/sqrt(1.*kxm%n)
            tfield(:,inode)=ws(:,inode); call sfftw_execute(planF(inode)); ws(:,inode)=tfield(:,inode)/sqrt(1.*kxm%n)
            gr0(:,inode)=wr(:,inode); gs0(:,inode)=ws(:,inode);

!      Operations in the wavenumber domain
!      phase shift
                do ikx = 1,kxm%n
                   sqa = rec%w%all(iw)**2 / vwbott**2 - kxm%all(ikx)**2

                   if (sqa > 0) then
                       gs0(ikx,inode)=gr0(ikx,inode)*exp(cmplx(0, dwbott*sqrt(sqa)))
                       gr0(ikx,inode)=gs0(ikx,inode)*exp(cmplx(0,-dwbott*sqrt(sqa)))
                   else
                       gs0(ikx,inode)=gr0(ikx,inode)*exp(-sqrt(-sqa)*dwbott)
                       gr0(ikx,inode)=gs0(ikx,inode)*exp(-sqrt(-sqa)*dwbott)
                   end if 
                end do 

                tfield(:,inode)=gr0(:,inode); call sfftw_execute(planI(inode)); wr(:,inode)=tfield(:,inode)/sqrt(1.*kxm%n)
                tfield(:,inode)=gs0(:,inode); call sfftw_execute(planI(inode)); ws(:,inode)=tfield(:,inode)/sqrt(1.*kxm%n)

        end do 
    
        recv(:,iw,inode)=wr(:,inode)
        wav(:,iw,inode)=ws(:,inode)

    end do
    end if
!
!-------------------------------------------------------
    omega:do iw = 1,rec%w%n

        if (allocated(ws)) deallocate(ws);	allocate(ws(rec%h%n,nnode));			ws=0
        if (allocated(wr)) deallocate(wr);	allocate(wr(rec%h%n,nnode));			wr=0
        if (allocated(gs)) deallocate(gs);	allocate(gs(rec%h%n,nref,nnode));		gs=0
        if (allocated(gr)) deallocate(gr);	allocate(gr(rec%h%n,nref,nnode));		gr=0
        if (allocated(gs0)) deallocate(gs0);	allocate(gs0(rec%h%n,nnode));			gs0=0
        if (allocated(gr0)) deallocate(gr0);	allocate(gr0(rec%h%n,nnode));			gr0=0
        if (allocated(sqa0)) deallocate(sqa0);	allocate(sqa0(rec%h%n,nnode));			sqa0=0

        ws(:,inode)=wav(:,iw,inode) ! to be modified to handle generalized sources
        wr(:,inode)=recv(:,iw,inode)

        depth:do iz=nwbott+1,img%z%n

            tfield(:,inode)=ws(:,inode); call sfftw_execute(planF(inode)); ws(:,inode)=tfield(:,inode)/sqrt(1.*kxm%n)
            tfield(:,inode)=wr(:,inode); call sfftw_execute(planF(inode)); wr(:,inode)=tfield(:,inode)/sqrt(1.*kxm%n)

            minv=minval(vel(iz,:,inode)); maxv=maxval(vel(iz,:,inode))

            vxz:if (minv .eq. maxv) then
                gs0(:,inode)=ws(:,inode);  gr0(:,inode)=wr(:,inode);
!      Operations in the wavenumber domain
!      phase shift
        if (mod(iw,10)==1) write(2,*)  "2 ",iw,iz,maxval(cabs(gs0)),maxval(cabs(gr0))
                   sqa0(:,inode) = rec%w%all(iw)**2 / minv**2 - kxm%all**2
                   evanescent1:where (sqa0(:,inode) > 0)
                       gs0(:,inode)=gs0(:,inode)*exp(cmplx(0, v%z%d*sqrt(sqa0(:,inode))))
                       gr0(:,inode)=gr0(:,inode)*exp(cmplx(0,-v%z%d*sqrt(sqa0(:,inode))))
                   elsewhere
                       gs0(:,inode)=gs0(:,inode)*exp(-sqrt(-sqa0(:,inode))*v%z%d)
                       gr0(:,inode)=gr0(:,inode)*exp(-sqrt(-sqa0(:,inode))*v%z%d)
                   end where evanescent1

! Operations in x            

                   tfield(:,inode)=gs0(:,inode); call sfftw_execute(planI(inode)); ws(:,inode)=tfield(:,inode)/sqrt(1.*kxm%n)
                   tfield(:,inode)=gr0(:,inode); call sfftw_execute(planI(inode)); wr(:,inode)=tfield(:,inode)/sqrt(1.*kxm%n)

            else;                                                                                       ! v(z) condition
                velref: do iref=1,nref

                    gs(:,iref,inode)=ws(:,inode);  gr(:,iref,inode)=wr(:,inode); 
!      phase shift
                      sqa0(:,inode) = rec%w%all(iw)**2 / vref(iref,iz,ixs)**2 - kxm%all**2;
                      evanescent2:where (sqa0(:,inode) > 0)
                          gs(:,iref,inode)=gs(:,iref,inode)*exp(cmplx(0, v%z%d*sqrt(sqa0(:,inode))))
                          gr(:,iref,inode)=gr(:,iref,inode)*exp(cmplx(0, -v%z%d*sqrt(sqa0(:,inode))))
                      elsewhere
                          gs(:,iref,inode)=gs(:,iref,inode)*exp(-sqrt(-sqa0(:,inode))*v%z%d)
                          gr(:,iref,inode)=gr(:,iref,inode)*exp(-sqrt(-sqa0(:,inode))*v%z%d)
                      end where evanescent2

! Operations in x            

                      tfield(:,inode)=gs(:,iref,inode); call sfftw_execute(planI(inode)); gs(:,iref,inode)=tfield(:,inode)/sqrt(1.*kxm%n)
                      tfield(:,inode)=gr(:,iref,inode); call sfftw_execute(planI(inode)); gr(:,iref,inode)=tfield(:,inode)/sqrt(1.*kxm%n)

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

                do iref=1,nref-1
                interpwvfld: where (vel(iz,:,inode) >= vref(iref,iz,ixs) .and. vel(iz,:,inode) < vref(iref+1,iz,ixs))
                     ws(:,inode)=gs(:,iref,inode)
                     wr(:,inode)=gr(:,iref,inode)
                     end where interpwvfld
                end do
                where (vel(iz,:,inode) >= vref(nref,iz,ixs))
                      ws(:,inode)=gs(:,nref,inode)
                      wr(:,inode)=gr(:,nref,inode)
                end where 

!                do ix=1,kxm%n
!                    do iref=1,nref-1
!                        if (vref(iref,iz,ixs) .ne. vref(iref+1,iz,ixs)) then
!                            if (vel(iz,ix,inode) >= vref(iref,iz,ixs) .and. vel(iz,ix,inode) < vref(iref+1,iz,ixs)) then
!                                f=(vel(iz,ix,inode) - vref(iref,iz,ixs))/vel(iz,ix,inode)
!                                g=1-f
!                                ws(ix,inode)=g*gs(ix,iref,inode)+f*gs(ix,iref+1,inode)
!                                wr(ix,inode)=g*gr(ix,iref,inode)+f*gr(ix,iref+1,inode)
!!                                ws(ix,inode)=gs(ix,iref,inode)
!!                                wr(ix,inode)=gr(ix,iref,inode)
!                            elseif (vel(iz,ix,inode) >= vref(nref,iz,ixs)) then
!                                ws(ix,inode)=gs(ix,nref,inode)
!                                wr(ix,inode)=gr(ix,nref,inode)
!                            end if
!                        else
!                            ws(ix,inode)=gs(ix,iref,inode)
!                            wr(ix,inode)=gr(ix,iref,inode)
!                        end if                            
!                    end do
!                end do
           end if vxz

!
! Imaging Condition
!
            do ish=1,img%sh%n
                ihso=nint(((ish-1)*img%sh%d+img%sh%o)/(img%xm%d))
                do ixw=1,kxm%n
                    if (ixw-ihso > 1 .and. ixw-ihso < kxm%n .and. ixw+ihso > 1 .and. ixw+ihso < kxm%n) then
                         GG(iz,ish,ixw,inode)=GG(iz,ish,ixw,inode)+conjg(ws(ixw-ihso,inode))*(wr(ixw+ihso,inode))             ! correlation to image
                    end if
                end do
            end do
        end do depth
!        if (mod(iw,10)==1) write(0,4)  "  DONE FREQUENCY # ",iw," OUT OF ",rec%w%n," / SHOT # ",ixs
4   format(a19,i3,a8,i3,a10,i3)
    end do omega

!    call date_and_time(values=time1)
!    finish_time = time1 (5) * 3600 + time1 (6) * 60 + time1 (7) + 0.001 * time1 (8)
    write(0,*) " DONE SHOT # ",ixs," OUT OF ",rec%x%n!," in ", (finish_time - start_time)/60," (min)"
!5   format(a13,i3,a8,i3,a4,f6.2,a6)

write(0,*) "Imaging for shot # ",ixs," is done in node: ",inode

!$OMP CRITICAL
     image(:,:,ni:ni+kxm%n-1)=image(:,:,ni:ni+kxm%n-1)+GG(:,:,:,inode);
!$OMP END CRITICAL

write(0,*) "After Critical",inode

end do
!$OMP end do
!$OMP end parallel

end subroutine
end module
