module Shan_post3dvz_mod

  use ooio, only: axis
!  use ooio
  use fft_data_mod, only: fft_type, init_fft_1d, init_rfft_1d, fft_1d_data, rfft_1d_data, clean_fft, clean_rfft
  implicit none

  private
  public :: post3dvz_init, post3dvz_mixdom, clean_post3dvz, post3dvz_fd,post3dvz_ffd
  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            :: nx, nt
  real               :: dx, dz
  type(fft_type)     :: fft_1d_id, rft_1d_id
  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 post3dvz_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
    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 post3dvz_init
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine post3dvz_mixdom(C, v)
    ! splitstep downward continuation

    complex, dimension (:,:), intent(inout) :: C !wavefield in (w,x) domain
    real   , dimension (nx)  , intent(in)    :: v ! 0.5*interval velocity

    real                                    :: v0! reference velocity
    real                                    :: sqa !sqrt(kz)
    integer                                 :: iw,ik

    v0 =sum(v)/k%n

    freq:do iw = 1, w%n
       call fft_1d_data(fft_1d_id, -1, k%n, C(iw,:))

       !Operations in the wavenumber domain
       splitstep_wvnmbr_dom:do ik = 1, k%n
          !---phase shift
          sqa = w%all(iw)**2 / v0**2 - k%all(ik)**2
          avoid_evan:if ( sqa > 0 ) then
             C(iw,ik) = C(iw,ik) * exp( i * sign * dz * sqrt(sqa) )
          else
             C(iw,ik) = 0.
          end if avoid_evan
       end do splitstep_wvnmbr_dom

       call fft_1d_data(fft_1d_id, 1, k%n, C(iw,:))

       !---split step
       C(iw,:) = C(iw,:) * exp(i*sign*dz*w%all(iw)*(1./v-1./v0))

    end do freq
  end subroutine post3dvz_mixdom
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine clean_post3dvz()

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

  end subroutine clean_post3dvz
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  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
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

subroutine post3dvz_fd(C, v)
    ! Implicit finite-difference downward continuation

    complex, dimension (:,:), intent(inout) :: C !wavefield in (w,x) domain
    real   , dimension (nx)  , intent(in)    :: v ! 0.5*interval velocity
    integer :: iw
    freq:do iw = 1, w%n
      C(iw,:) = C(iw,:) * exp(i*sign*dz*w%all(iw)*(1./v))
      call fd(C(iw,:),v,iw)
    end do freq

end subroutine post3dvz_fd


subroutine post3dvz_ffd(C, v)
    ! Fourier finite-difference  downward continuation

    complex, dimension (:,:), intent(inout) :: C !wavefield in (w,x) domain
    real   , dimension (nx)  , intent(in)    :: v ! 0.5*interval velocity

    real                                    :: v0! reference velocity
    real                                    :: sqa !sqrt(kz)
    integer                                 :: iw,ik

    v0 =minval(v)

    freq:do iw = 1, w%n
       call fft_1d_data(fft_1d_id, -1, k%n, C(iw,:))

       !Operations in the wavenumber domain
       splitstep_wvnmbr_dom:do ik = 1, k%n
          !---phase shift
          sqa = w%all(iw)**2 / v0**2 - k%all(ik)**2
          avoid_evan:if ( sqa > 0 ) then
             C(iw,ik) = C(iw,ik) * exp( i * sign * dz * sqrt(sqa) )
          else
             C(iw,ik) = 0.
          end if avoid_evan
       end do splitstep_wvnmbr_dom
       call fft_1d_data(fft_1d_id, 1, k%n, C(iw,:))
       !---split step
       C(iw,:) = C(iw,:) * exp(i*sign*dz*w%all(iw)*(1./v-1./v0))
       call fd_forffd(C(iw,:),v,v0,iw) 

    end do freq




end subroutine post3dvz_ffd




subroutine fd_forffd(u,v,v0,iw)
complex :: u(:)
real :: v(:)
real :: v0
integer :: iw

real :: tx,a,b
complex,allocatable :: uj(:),ra(:),rb(:),rc(:)
integer :: ix
complex :: url,urr,ror,rol,dd
real :: vw,v0w

a=0.5
b=0.25
!a=0.4765
!b=0.3725
tx=0.11

allocate(uj(nx))
allocate(ra(nx))
allocate(rb(nx))
allocate(rc(nx))
v0w=v0/(w%all(iw))
do ix=1,nx
  uj(ix)=u(ix)
end do
do ix=1,nx
  if (ix==1)  then
    url=0.0
  else
    url=uj(ix-1)
  end if
  if (ix==nx) then
    urr=0.0
  else
    urr=uj(ix+1)
  end if
  vw=v(ix)/(w%all(iw))
  ror=cmplx(tx+b*(vw*vw+v0w*vw+v0w*v0w)/(dx*dx),sign*a*(vw-v0w)*dz/(2.0*dx*dx)) 
  rol=conjg(ror)
  u(ix)=uj(ix) +ror*(url+urr-2.0*uj(ix))
  ra(ix)=1.0-2.0*rol
  if (ix <nx) rb(ix) =rol
  if (ix >1) rc(ix-1)=rol
end do
call thr(ra,rb,rc,u)
deallocate(uj,ra,rb,rc)
end subroutine fd_forffd




subroutine fd(u,v,iw)
complex :: u(:)
real :: v(:)
integer :: iw

real :: tx,a,b
complex,allocatable :: uj(:),ra(:),rb(:),rc(:)
integer :: ix
complex :: url,urr,ror,rol,dd
real :: vw

a=0.5
b=0.25
b=0.0
!a=0.4765
!b=0.3725
tx=0.11

allocate(uj(nx))
allocate(ra(nx))
allocate(rb(nx))
allocate(rc(nx))
do ix=1,nx
  uj(ix)=u(ix)
end do
do ix=1,nx
  if (ix==1)  then
    url=0.0
  else
    url=uj(ix-1)
  end if
  if (ix==nx) then
    urr=0.0
  else
    urr=uj(ix+1)
  end if
  vw=v(ix)/(w%all(iw))
  ror=cmplx(tx+b*vw*vw/(dx*dx),sign*a*vw*dz/(2.0*dx*dx)) 
  rol=conjg(ror)
  u(ix)=uj(ix) +ror*(url+urr-2.0*uj(ix))
  ra(ix)=1.0-2.0*rol
  if (ix <nx) rb(ix) =rol
  if (ix >1) rc(ix-1)=rol
end do
call thr(ra,rb,rc,u)
deallocate(uj,ra,rb,rc)
end subroutine fd


!tri-diagonal solver
subroutine thr(a,b,c,u)
complex :: a(:),b(:),c(:),u(:)
integer ::ix
b(1)=b(1)/a(1)
u(1)=u(1)/a(1)
do ix=2,nx
   a(ix)=a(ix)-b(ix-1)*c(ix-1)
   u(ix)=(u(ix)-c(ix-1)*u(ix-1))/a(ix)
   if(ix.le.nx-1) then
      b(ix)=b(ix)/a(ix)
   endif
enddo
do ix=nx-1,1,-1
  u(ix)=u(ix)-b(ix)*u(ix+1)
enddo
end subroutine thr





end module Shan_post3dvz_mod
