module propaga_mod

  use fft_data_mod
  implicit none

  private
  public :: prop_init, prop_phsh, clean_prop_phsh
  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
  real, dimension(:), allocatable	:: 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(nz_in, dz_in, oz_in, nxm_in, dxm_in, oxm_in, nx_in, dx_in, ox_in, nt_in, dt_in, ot_in) result(stat)

    integer                :: iw, ik
    integer		   :: stat
    integer                :: nz_in, nxm_in, nx_in, nt_in
    real 		   :: dz_in, oz_in, dxm_in, oxm_in, dx_in, ox_in, dt_in, ot_in

    sign=-1 ! if(dowcon_as_shot) sign = +1
    nz=nz_in;    dz=dz_in;    oz=oz_in
    nxm=nxm_in;  dxm=dxm_in;  oxm=oxm_in
    nx=nx_in;    dx=dx_in;    ox=ox_in
    nt=nt_in;    dt=dt_in;    ot=ot_in

    write(0,*) 'nz  = ',nz,'dz  = ',dz,'oz  = ',oz
    write(0,*) 'nxm = ',nxm,'dxm = ',dxm,'oxm = ',oxm
    write(0,*) 'nt  = ',nt,'dt  = ',dt,'ot  = ',ot
    write(0,*) 'nx  = ',nx,'dx  = ',dx,'ox  = ',ox

    rft_1d_id = init_rfft_1d(nt)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    nw = nt/2
    dw = 1./(nt*dt)
    ow = 0.! dw				 ! no zero frequency component
    allocate(w(nw)); w=0.
    do iw = 1, nw
       w(iw) = ow  + (iw - 1) * dw
    end do
    w = 2*pi * w
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    fft_1d_id = init_fft_1d(nx)
    nk = nx
    allocate(k(nk)); k=0.

    dk = 1./(nk*dx)
    ok = 1./(2*dx)-dk			!k does not increase uniformly!
    do ik=1, nk
       k(ik) = mod(ok+(ik-1)*dk,1./dx)-ok
    end do
    k = 2*pi * k
    
    stat = 0
   
  end function prop_init
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine prop_phsh(C, v, D)
! phase-shift modeling

    complex, dimension (:,:)    	       :: C !model in (z,x) domain
    complex, dimension (:,:), allocatable      :: U !wavefield in (w,kx) domain and temporary A
    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(nw,nk))
    U = 0
    do iz = 1,nz
       call fft_1d_data(fft_1d_id, -1, nk, C(iz,:))
    end do
    
!---frequency loop
    do iw = 1, nw
       write(0,5) (iw-1)*100./nw, '% modeled'
 5     format(f5.1,a10)       

!Operations in the wavenumber domain
       do ik = 1, nk
! z
          do iz = nz,1,-1
             v0 = sum(v(iz,:))/nk

!---phase shift term
             sqa = w(iw)**2 / v0**2 - k(ik)**2
!---avoid evan
             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,nw
       call fft_1d_data(fft_1d_id, 1, nk, U(iw,:))
    end do

    D = real(rft_bck_axis1(U))

  end subroutine prop_phsh
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine clean_prop_phsh()

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

  end subroutine clean_prop_phsh
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  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 propaga_mod
