module grf_vz_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, nz
  real               :: dx, dz, dt
  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
    dt = t%d
    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 post3dvz_init
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine post3dvz_mixdom(C, v, D)
    ! splitstep modeling

    complex, dimension (:,:)    	       :: C !wavefield in (w,x) domain
    real, dimension (:,:,:),  intent(out)      :: D !wavefield in (w,kx) domain
    real   , dimension (nz,nx),  intent(in)    :: v ! 0.5*interval velocity
    real   , dimension (k%n)  		       :: v1 ! 0.5*interval velocity

    real                                       :: v0! reference velocity
    real                                       :: sqa !sqrt(kz)
    integer                                    :: iw,ik,indw,it,ix,iz
    complex				       :: pow
    complex, dimension (nx)	    	       :: pow1 !


    D = 0
!---frequency loop
!    do iw = 1, w%n
    iw=20
       do iz = 1,nz
!          write(0,*) iz*100./nz, '% modeled'
          v0 = sum(v(iz,:))/k%n
	  write(0,*) sum(1./v(iz,:)-1./v0) 

!---split step
	  pow1 = exp(i*sign*dz*w%all(iw)*(1./v(iz,:)-1./v0))
          C(iw,:) = C(iw,:) * pow1
	  
          call fft_1d_data(fft_1d_id, -1, k%n, C(iw,:))

!Operations in the wavenumber domain
          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 * dz * sqrt(sqa) )
             else
               C(iw,ik) = 0.
             end if
          end do

          call fft_1d_data(fft_1d_id, 1, k%n, C(iw,:))
       
!---evolution in time
          do it=1,nt
	     pow = exp( -i * w%all(iw) * it * dt )
	     do ix=1,nx
		D(iz,ix,it) = D(iz,ix,it) + C(iw,ix) * pow
             end do
          end do


!       D(iw,:) = C(iw,:)
!       call fft_1d_data(fft_1d_id, -1, k%n, D(iw,:))

!       end do 

    end do 

  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(P, v, ang)
    ! Implicit finite-difference downward continuation

    complex, dimension (:,:), intent(inout)	:: P !wavefield in (w,x) domain
    real   , dimension (nx)  , intent(in)	:: v ! 0.5*interval velocity
    integer 					:: iw,ix
    real 					:: ang, coef1, coef2
    real					:: vs
    complex, dimension(nx)			:: alfa, beta, a, b, c, U 

   freq:do iw = 1, w%n
         eq: if (ang.eq.15.) then !---'finite-difference 15deg'
	     coef1=0.5;coef2=0.
         elseif (ang.eq.45.) then !---'finite-difference 45deg'
	     coef1=0.5;coef2=0.25
	 elseif (ang.eq.60.) then !---'finite-difference 60deg'
	     coef1=0.47824206;coef2=0.37636953
	 end if eq

         alfa = coef1*v*dz/(2.*(0.,1.)*w%all(iw)*dx**2);beta=coef2*(v/(w%all(iw)*dx))**2
         a = -(alfa-beta); b = (1.+2.*(alfa-beta)); c(1:nx-1) = -(alfa(2:nx)-beta(2:nx))  ! tridiagonal matrix coeffients

         difterm:do ix=2,nx-1
                U(ix) = (alfa(ix)+beta(ix))*P(iw,ix+1)+(1.-2.*(alfa(ix)+beta(ix)))*P(iw,ix)+(alfa(ix)+beta(ix))*P(iw,ix-1)
         end do difterm

	 U(1)=0.
	 U(nx)=0.

         call thr(b,a,c,U)

         P(iw,:) = U*exp(i*sign*dz*w%all(iw)*(1/v))
         
    end do freq

end subroutine post3dvz_fd
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

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

    complex, dimension (:,:), intent(inout)	:: P !wavefield in (w,x) domain
    real   , dimension (nx)  , intent(in)	:: v ! 0.5*interval velocity
    integer 					:: iw,ix,ik
    real 					:: v0,sqa,ang,coef1
    real, dimension(nx)				:: coef2
    real , dimension(nx)			:: gama
    complex, dimension(nx)			:: alfa,beta, a, b, c, U

    v0 = minval(v)

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

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

         eq1: if (ang.eq.15.) then !---'finite-difference 15deg'
	     coef1=0.5;coef2=0.
         elseif (ang.eq.45.) then !---'finite-difference 45deg'
	     coef1=0.5;coef2=0.25*((v0/v)**2+v0/v+1)
	 end if eq1

	 gama=(1-v0/v)

         alfa = gama*coef1*v*dz/(2.*(0.,1.)*w%all(iw)*dx**2)
	 beta = coef2*(v/(w%all(iw)*dx))**2
         a = -(alfa-beta); b = (1+2.*(alfa-beta)); c(1:nx-1) = -(alfa(2:nx)-beta(2:nx))  ! tridiagonal matrix coeffients

         difterm1:do ix=2,nx-1
	        U(ix) = (alfa(ix)+beta(ix))*P(iw,ix+1)+(1-2.*(alfa(ix)+beta(ix)))*P(iw,ix)+(alfa(ix)+beta(ix))*P(iw,ix-1)
         end do difterm1
         U(1)=0.
         U(nx)=0.
         call thr(b,a,c,U)
         P(iw,:) = U

    end do freq1    
end subroutine post3dvz_ffd


!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 grf_vz_mod
