module OW_pspi

use OWvel
use OW_type
use OW_image
use OW_fftw
use wei_myfft_mod
use OW_parms
use OW_imageIO
use OW_exec_time

implicit none

public						:: owpspi_init, owpspi
private						:: isp,ixs

integer						:: isp,ixs
real						:: wb

contains

!---------------------------------------------------------
subroutine owpspi_init(isp_in,ixs_in,wb_in)

integer						:: isp_in,ixs_in
real						:: wb_in

isp=isp_in
ixs=ixs_in
wb=wb_in

end subroutine
!---------------------------------------------------------
subroutine owpspi(recv,source,vel)

integer						:: iw,iz,ikx,iref,ixw,ish,ino
integer						:: nv,ihso,nwb,node,inode
integer,allocatable					:: indref(:)
real						:: sqa,f,g,start,finish
real						:: minv, maxv,vel(:,:)
real, allocatable					:: vref(:),sqa0(:)
complex, allocatable				:: ws(:),wr(:),gs0(:),gr0(:),gs(:,:),gr(:,:),tfld(:),a(:),b(:)
!real, allocatable					:: vref(:,:),sqa0(:,:)
!complex, allocatable				:: ws(:,:),wr(:,:),gs0(:,:),gr0(:,:),gs(:,:,:),gr(:,:,:),tfld(:,:)
complex						:: recv(:,:),source(:,:)

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

real,allocatable					:: image(:,:,:),GG(:,:,:)

node=omp_get_num_threads()
node=8

allocate(GG(kxm%n,img%sh%n,img%z%n))
allocate(image(kxm%n,img%sh%n,img%z%n))

!--------------------------------------------------
nwb=floor(wb/img%z%d)+1

!--------------------------------------------------
! Setting FFTW plans
allocate(tfld(kxm%n));tfld=0
call owfftw(tfld)
deallocate(tfld)

! $OMP parallel do private (inode,&
!                             iw,&
!                             iz,&
!                            ikx,&
!                           iref,&
!                            ish,&
!                           ihso,&
!                            ixw,&
!                             wr,&
!                             gr,&
!                            gr0,&
!                             ws,&
!                             gs,&
!                            gs0,&
!                           vref,&
!                           sqa0,&
!                            sqa,&
!                           minv,&
!                           maxv,&
!                           tfld,&
!                         indref,&
!                              f,&
!                              g) schedule(dynamic,1) num_threads(8)
! !$OMP DO schedule(dynamic,1) 

omega:do iw = 1,rec%w%n

        if(allocated(ws)) deallocate(ws); allocate(ws(kxm%n)); ws=0.
        if(allocated(wr)) deallocate(wr); allocate(wr(kxm%n)); wr=0.

        inode=omp_get_thread_num()+1; 

        ws=source(:,iw) 
        wr=recv(:,iw) 

        write(0,*) cabs(sum(ws)),cabs(sum(wr))

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

            if(allocated(tfld)) deallocate(tfld); allocate(tfld(kxm%n));tfld=0

!!--------------------------------------------------
!! Compute vref
!!
!start=0.;finish=0. 
!call owexec_time(start)

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

!if (iw==1) write(0,222) iz,vref
!222 format(i4,5f8.1)

!call owexec_time(finish)
!write(0,*) "Finished computing vref in ", finish-start," (s)"

!!---------------------------------------------------
!!
            tfld=ws; call sfftw_execute(planF); ws=tfld/sqrt(1.*kxm%n);
            tfld=wr; call sfftw_execute(planF); wr=tfld/sqrt(1.*kxm%n)

        write(2,*) "1st fourier",iz,cabs(sum(ws)),cabs(sum(wr))

!start=0.;finish=0. 
!call owexec_time(start)

!            tfld=ws; call FFT1DX(tfld,1); ws=tfld;
!            tfld=wr; call FFT1DX(tfld,1); wr=tfld;

!call owexec_time(finish)
!write(0,*) "Finished 1st Fourier transf in ", finish-start," (s)"

           vxz:if (minv .eq. maxv .and. vref(size(vref)).lt.vsal) then
               if(allocated(gs0)) deallocate(gs0); allocate(gs0(kxm%n));gs0=ws
               if(allocated(gr0)) deallocate(gr0); allocate(gr0(kxm%n));gr0=ws
!
!      Operations in the wavenumber domain
!      phase shift
!start=0.;finish=0. 
!call owexec_time(start)

                wavenumber1:do ikx = 1,kxm%n
                   sqa = rec%w%all(iw)**2 / minv**2 - kxm%all(ikx)**2

                   evanescent1:if (sqa > 0) then
                       gs0(ikx)=gs0(ikx)*exp(cmplx(0,v%z%d*sqrt(sqa)))
                       gr0(ikx)=gr0(ikx)*exp(cmplx(0,-v%z%d*sqrt(sqa)))
                   else
                       gs0(ikx)=gs0(ikx)*exp(-sqrt(-sqa)*v%z%d)
                       gr0(ikx)=gr0(ikx)*exp(-sqrt(-sqa)*v%z%d)
                   end if evanescent1
                end do wavenumber1


!call owexec_time(finish)
!write(0,*) "Finished propagating cte.vel. in ", finish-start," (s)"

!start=0.;finish=0. 
!call owexec_time(start)

!            tfld=gs0; call FFT1DX(tfld,-1); ws=tfld;
!            tfld=gr0; call FFT1DX(tfld,-1); wr=tfld;

!call owexec_time(finish)
!write(0,*) "Finished 2nd Fourier transf in ", finish-start," (s)"

             tfld=gs0; call sfftw_execute(planI); ws=tfld/sqrt(1.*kxm%n)
             tfld=gr0; call sfftw_execute(planI); wr=tfld/sqrt(1.*kxm%n)

        write(2,*) "2nd fourier",iz,cabs(sum(ws)),cabs(sum(wr))

!
            else

                if (allocated(a).and.allocated(b)) deallocate(a,b); allocate(a(kxm%n),b(kxm%n))
                a=ws;b=wr
                velref: do iref=1,size(vref)                 

                    gs(:,iref)=ws
                    gr(:,iref)=wr
                    sqa=0

        write(2,*) "g = w / vref",iz,sum(cabs(gs(:,iref))),sum(cabs(gr(:,iref)))

!start=0.;finish=0. 
!call owexec_time(start)

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

        write(2,*) "phase shift vref",iz,sum(cabs(gs(:,iref))),sum(cabs(gr(:,iref)))


!call owexec_time(finish)
!write(0,*) "Finished propagating w/ref.vel. in ", finish-start," (s)"

! Operations in x            

!start=0.;finish=0. 
!call owexec_time(start)

!            tfld=gs(:,iref); call FFT1DX(tfld,-1); gs(:,iref)=tfld;
!            tfld=gr(:,iref); call FFT1DX(tfld,-1); gr(:,iref)=tfld;

!call owexec_time(finish)
!write(0,*) "Finished 3rd Fourier transf in ", finish-start," (s)"

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

!      split step

!start=0.;finish=0. 
!call owexec_time(start)

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

        write(2,*) "split step vref",iz,sum(cabs(gs(:,iref))),sum(cabs(gr(:,iref)))


!call owexec_time(finish)
!write(0,*) "Finished split step in ", finish-start," (s)"

                end do velref

!start=0.;finish=0. 
!call owexec_time(start)

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

        write(2,*) "After interpolation ",iz,cabs(sum(ws)),cabs(sum(wr))


!call owexec_time(finish)
!write(0,*) "Finished interpolating wvld in ", finish-start," (s)"

!
! Imaging Condition
!

!start=0.;finish=0. 
!call owexec_time(start)

           do ish=1,img%sh%n
               ihso=floor(((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(ixw,ish,iz)=GG(ixw,ish,iz)+real(conjg(ws(ixw-ihso))*(wr(ixw+ihso)))             ! correlation to image
                   end if
               end do
           end do

        write(2,*) "After imaging ",iz,sum(GG(:,1,iz))



        end do depth

!call owexec_time(finish)
!write(0,*) "Finished imaging in ", finish-start," (s)"

!
        ! $OMP CRITICAL
           image=image+GG

!start=0.;finish=0. 
!call owexec_time(start)

           call owimage_io(image,ixs)

!call owexec_time(finish)
!write(0,*) "Finished writing image in ", finish-start," (s)"

        ! $OMP END CRITICAL
!          if (mod(iw,20)==1) write(0,5)  "DONE FREQUENCY # ",iw," OUT OF ",rec%w%n," / SHOT # ",ixs
          write(0,5)  "DONE FREQUENCY # ",iw," OUT OF ",rec%w%n," / SHOT # ",ixs
5                            format(A17,i5,2(A8,i5))

end do omega

! $OMP END PARALLEL DO

end subroutine

end module
