module propaga2_mod

  use fft_data_mod
  use ooio

  implicit none

  private
  public :: prop_init, prop_phsh, clean_prop, prop_pspi, prop_splstep
  public :: rft_fwd_axis1, rft_bck_axis1

  real,    parameter 			:: pi=3.1415926535898
  complex, parameter 			:: i = cmplx(0.,1.)
  integer            			:: sign !sign of phaseshift in exponential
  integer, private            		:: nz, nxm, nx, nt, nk, nw
  real, private        			:: dz, oz, dxm, oxm, dx, ox, dt, ot, dk, ok, dw, ow
  type(fft_type)     			:: fft_1d_id, rft_1d_id, fft_2d_id, fft_1d_idw
  type(axis) 				:: w, k

  ! Bug danger! The common variables should be modified ONLY by init & clean. Or at worst, by a single subroutine only.

contains
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  function prop_init(x, z, t) result(omega)

    type(axis), intent(in) :: x, z, t
    integer                :: iw, ik
    type(axis)             :: omega

    sign=-1 ! if(dowcon_as_shot) sign = +1
    nx = x%n
    dx = x%d
    nt = t%n
    dz = z%d
    nz = z%n
    rft_1d_id = init_rfft_1d(nt)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    w%n = t%n/2
    w%d = 1./(t%n*t%d)
    w%o = w%d ! no zero frequency component
    allocate(w%all(w%n)); w%all=0.
    do iw = 1, w%n
       w%all(iw) = w%o  + (iw - 1) * w%d
    end do
    w%all = 2*pi * w%all

    omega = w
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    fft_1d_id = init_fft_1d(nx)
    k%n = nx
    allocate(k%all(k%n)); k%all=0.

    k%d = 1./(k%n*dx)
    k%o = 1./(2*dx)-k%d !k does not increase uniformly!
    do ik=1, k%n
       k%all(ik) = mod(k%o+(ik-1)*k%d,1./dx)-k%o
    end do

    k%all = 2*pi * k%all

  end function prop_init
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine prop_phsh(M, v, D)
! phase-shift modeling

    real, dimension (:,:)  	  	       :: M !model in (z,x) domain
    complex, dimension (:,:), allocatable      :: C !model in (z,x) domain
    complex, dimension (:,:), allocatable      :: U !wavefield in (w,kx) domain
    real, dimension (:,:)		       :: D !wavefield in (t,x) domain
    real   , dimension (nz,nx),  intent(in)    :: v ! 0.5*interval velocity
    real                                       :: v0! reference velocity
    real                                       :: sqa,dza !sqrt(kz)
    integer                                    :: iw,ik,indw,it,ix,iz

    allocate(U(w%n,k%n)); U = 0
    allocate(C(nz,nx)); C = 0

!---frequency loop
   do iw = 1, w%n
         if (mod(iw,10).eq.1) write(0,5) (iw-1)*100./w%n, '% modeled'
5        format(f5.1,a10)       
      C = cmplx(M)
! z
      do iz = nz,1,-1
         v0 = sum(v(iz,:))/k%n
         call fft_1d_data(fft_1d_id, -1, k%n, C(iz,:))
!Operations in the wavenumber domain
         do ik = 1, k%n

!---phase shift term
            sqa = w%all(iw)**2 / v0**2 - k%all(ik)**2
!---avoid evanescent energy
            if ( sqa > 0 ) then
              U(iw,ik) = U(iw,ik) * exp( i * dz * sqrt(sqa) ) + C(iz,ik)
            else
              U(iw,ik) = 0.
            end if
         end do
      end do 
      call fft_1d_data(fft_1d_id, 1, k%n, U(iw,:))
   end do 

   D = rft_bck_axis1(U)

  end subroutine prop_phsh

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine prop_pspi(nv0, M, v, D)
! phase-shift plus interpolation modeling

    real, 	dimension (:,:)  	  	:: M !model in (z,x) domain
    complex,	dimension (:,:), allocatable	:: C !model in (z,x) domain
    complex, 	dimension (:,:), allocatable	:: U !wavefield in (w,kx) domain
    real, 	dimension (:,:)			:: D !wavefield in (t,x) domain
    complex, 	dimension (:,:,:), allocatable	:: U1,U2 !wavefields in (w,kx,v) domain
    real, 	dimension (:,:), intent(in)	:: v ! 0.5*interval velocity
    integer                             	:: iw,ik,indf,indg,it,ix,iz,nv0,iv
    real					:: v0, sqa ! reference velocity & sqrt(kz)
    real					:: v0min,v0max,dv,f,g,v1
    real					:: fi1,fi2,amp1,amp2,phu,ampu,ru,imu
!
    allocate(U1(w%n,nx,nv0),U2(w%n,nx,nv0))
!
    allocate(U(w%n,k%n)); U = 0
    allocate(C(nz,nx)); C = 0
!
!---z
   do iz = nz,1,-1
      C = cmplx(M)
      if (mod(iz,10).eq.1) write(0,5) (iz-1)*100./nz, '% modeled'
5     format(f5.1,a10)       
      call fft_1d_data(fft_1d_id, -1, k%n, C(iz,:))
!---definition of reference velocities
      v0min=minval(v(iz,:)); v0max=maxval(v(iz,:)); dv=(v0max-v0min)/(nv0-1);
!---loop in frequency 
      do iw = 1, w%n
!---loop in v0
         do iv=1,nv0
            v0=(iv-1)*dv+v0min
!---loop in kx
            do ik = 1, k%n
!---phase shift term
               sqa = w%all(iw)**2 / v0**2 - k%all(ik)**2

!---avoid evanescent energy
               if ( sqa > 0 ) then
!                 U1(iw,ik,iv) = U1(iw,ik,iv) * exp( i * dz * (sqrt(sqa) + w%all(iw)/v0) ) + C(iz,ik)
                 U1(iw,ik,iv) = U1(iw,ik,iv) * exp( i * dz * sqrt(sqa) ) + C(iz,ik)
               else
                 U1(iw,ik,iv) = 0.
               end if
            end do								! loop kx
	    U(iw,:)=U1(iw,:,iv)
	    call fft_1d_data(fft_1d_id, 1, k%n, U(iw,:))
	    U2(iw,:,iv)=U(iw,:)
         end do									! loop v
         U(iw,:) = 0
!---Linear interpolation
         do ix = 1,nx
	    do iv=1,nv0
               v0=(iv-1)*dv+v0min; v1=iv*dv+v0min
	       if (v1.eq.v0) then
	          U(iw,ix)=U2(iw,ix,iv)
	          exit
	       end if
	       f = (v(iz,ix)-v0)/dv; g = (v1-v(iz,ix))/dv
	       if (iv.lt.nv0) then
	          if ((v(iz,ix).le.v1).and.(v(iz,ix).ge.v0)) then
		     fi1=atan(aimag(U2(iw,ix,iv))/real(U2(iw,ix,iv)))
		     fi2=atan(aimag(U2(iw,ix,iv+1))/real(U2(iw,ix,iv+1)))
		     amp1=cabs(U2(iw,ix,iv))
		     amp2=cabs(U2(iw,ix,iv+1))
		     phu=g*fi1+f*fi2; ampu=g*amp1+f*amp2
		     ru=ampu/sqrt(1+tan(phu)**2);imu=ru*tan(phu)
		     U(iw,ix)=cmplx(ru,imu)
!	             U(iw,ix)=g*U2(iw,ix,iv) + f*U2(iw,ix,iv+1)
		     exit
	          end if
	       else
		  fi1=atan(aimag(U2(iw,ix,iv-1))/real(U2(iw,ix,iv-1)))
		  fi2=atan(aimag(U2(iw,ix,iv))/real(U2(iw,ix,iv)))
		  amp1=cabs(U2(iw,ix,iv-1))
		  amp2=cabs(U2(iw,ix,iv))
		  phu=g*fi1+f*fi2; ampu=g*amp1+f*amp2
		  ru=ampu/sqrt(1+tan(phu)**2);imu=ru*tan(phu)
		  U(iw,ix)=cmplx(ru,imu)
!	          U(iw,ix)=g*U2(iw,ix,iv-1) + f*U2(iw,ix,iv)
		  exit
	       end if		   
	    end do
	 end do									! linear interp.
!---
!      U(iw,:) = U(iw,:) * exp( - i * dz * w%all(iw)/v(iz,:) ) 
      end do 									! loop w
   end do 									! loop z
!
  
   D = rft_bck_axis1(U)

  end subroutine prop_pspi
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine prop_splstep(M, v, D)
!! slipt-step modeling

    real, dimension (:,:)  	  	       :: M !model in (z,x) domain
    complex, dimension (:,:), allocatable      :: C,A !model in (z,x) domain
    complex, dimension (:,:), allocatable      :: U !wavefield in (w,kx) domain
    real, dimension (:,:)		       :: D !wavefield in (t,x) domain
    real   , dimension (nz,nx),  intent(in)    :: v ! 0.5*interval velocity
    real                                       :: v0! reference velocity
    real                                       :: sqa,dza !sqrt(kz)
    integer                                    :: iw,ik,indw,it,ix,iz

    allocate(U(w%n,k%n)); U = 0
    allocate(C(nz,nx),A(nz,nx)); C = 0; A = 0;
    A = cmplx(M)
open(unit=2,file="lixo")
!---frequency loop
   do iw = 1, w%n
      if (mod(iw,10).eq.1) write(0,5) (iw-1)*100./w%n, '% modeled'
5     format(f5.1,a10)    
      C = A   
! z
      do iz = nz,1,-1
         v0 = sum(v(iz,:))/k%n
         C(iz,:) = C(iz,:) * exp(i*dz*w%all(iw)*(1./v(iz,:)-1./v0))
         call fft_1d_data(fft_1d_id, -1, k%n, C(iz,:))
!Operations in the wavenumber domain
         do ik = 1, k%n
!---phase shift term
            sqa = w%all(iw)**2 / v0**2 - k%all(ik)**2
!---avoid evanescent energy
            if ( sqa > 0 ) then
              U(iw,ik) = U(iw,ik) * exp( i * dz * sqrt(sqa) ) + C(iz,ik)
            else
              U(iw,ik) = 0.
            end if
         end do
      end do 
   end do 

   do iw = 1,w%n
      call fft_1d_data(fft_1d_id, 1, k%n, U(iw,:))
   end do

   D = rft_bck_axis1(U)
close(2)

end subroutine prop_splstep

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine clean_prop()

    call clean_rfft(rft_1d_id)
    call clean_fft (fft_1d_id)
    deallocate(k%all, w%all)

  end subroutine clean_prop
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  function rft_fwd_axis1(U)
    real   , dimension(:,:)        , intent(in) :: U
    complex, dimension(size(U,1)/2,  size(U,2)) :: rft_fwd_axis1
    complex, dimension(size(U,1)/2+1)           :: all_freq
    integer                                     :: ix

    do ix = 1, nx
       call rfft_1d_data(rft_1d_id, -1, nt, U(:,ix), all_freq)
       rft_fwd_axis1(:,ix) = all_freq(2:) ! Throw away the zero frequency
    end do
  end function rft_fwd_axis1
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  function rft_bck_axis1(C)
    complex, dimension(:,:), intent(in)        :: C
    real   , dimension(2*size(C,1), size(C,2)) :: rft_bck_axis1
    complex, dimension(1+size(C,1))            :: all_freq
    integer                                    :: ix

    all_freq (1)=0

    do ix = 1, nx
       all_freq(2:) = C(:,ix)
       call rfft_1d_data(rft_1d_id, 1, nt, rft_bck_axis1(:,ix), all_freq)
    end do
  end function rft_bck_axis1
end module propaga2_mod
