module post3dvz_mod

  use ooio, only: axis
  use base_slofft_mod
  implicit none

  private
  public :: post3dvz_init, post3dvz_mixdom, clean_post3dvz, post3dvz_fd, post3dvz_fd2, post3dvz_ffd, post3dvz_ffd2

  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(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

    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

    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,mode,vref)
    ! mode=1 Phase shift migration
    ! mode=2 split-step downward continuation

    complex, dimension (:,:), intent(inout) :: C !wavefield in (w,x) domain
    real   , dimension (nx) , intent(in)    :: v ! 0.5*interval velocity
    
    real                                    :: v0, vref! reference velocity
    real                                    :: sqa !sqrt(kz)
    integer                                 :: iw,ik
    integer                                 :: mode

    !!.. reference velocity

    if (mode .eq. 3) then
        v0=vref
    else
        v0 =sum(v)/k%n
    end if

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

       if ((mode .eq. 2) .or. (mode .eq. 3)) then
           !---split step
           C(iw,:) = C(iw,:) * exp(i*sign*dz*w%all(iw)*(1./v-1./v0))
       end if

    end do
  end subroutine post3dvz_mixdom
 

  subroutine clean_post3dvz()
    deallocate(k%all, w%all)
  end subroutine clean_post3dvz

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

    integer                                 :: iw, ix, mode
    complex, dimension (:,:), intent(inout) :: C !wavefield in (w,x) domain
    real   , dimension (:)  , intent(in)    :: v ! 0.5*interval velocity
    real                                    :: c1, c2, beta, vavg
    complex                                 :: alpha
    complex, dimension(:), allocatable      :: a, b, cc, r, u



    allocate(a(nx), b(nx), cc(nx), r(nx), u(nx));

!   add your implicit finite difference here 

 !   write(0,*) 'nw, nx', w%n, nx

    if (mode .eq. 1) then    ! Set coeff for 15 degree eqn
       
       c1=0.5; c2=0.0;
     
    elseif (mode .eq. 2) then   ! Set coeff for 45 degree eqn
    
       c1=0.5; c2=0.25;
    
    else                       ! Set coeff for 60 degree eqn

       c1=0.47824206; c2=0.37636953;

    endif
        

        do iw=1,w%n


           do ix=1,nx     ! Clear temp array
              a(ix)=0.; b(ix)=0.; cc(ix)=0.; 
              r(ix)=0.; u(ix)=0.; 
           enddo

           
           do ix=2,nx-1
   !           alpha= i*(c1+c2)*v(ix)*dz/(2.*w%all(iw)*dx**2);
              alpha= i*sign*(c1)*v(ix)*dz/(2.*w%all(iw)*dx**2);
              beta = c2*(v(ix)/w%all(iw)/dx)**2;


              a(ix)=beta-alpha;
              cc(ix)=beta-alpha;
              b(ix)=1-2.*beta+2.*alpha;

              r(ix) = (beta + alpha)*C(iw,ix-1) &
                   +(1 - 2.*beta - 2.*alpha)*C(iw,ix) & 
                   +(beta + alpha)*C(iw,ix+1);
           enddo

           a(1)=1.;  b(1)=1.;  cc(1)=1.;
           a(nx)=1.;  b(nx)=1.; cc(nx)=1.;

           r(1)=0.;   
           r(nx)=0.;

           ! Call tridiagonal solver
           call ctridag(a,b,cc,r,u,nx)

           ! Shifting term
           do ix=1,nx
   !            C(iw,ix) = C(iw,ix)*exp(i*sign*dz*w%all(iw)/v(ix));
              u(ix) = u(ix)*exp(i*sign*dz*w%all(iw)/v(ix));
           enddo

           C(iw,:)=u;

        enddo

        deallocate(a,b,cc,r,u);

  end subroutine 

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

    integer                                 :: ik, iw, ix, mode
    complex, dimension (:,:), intent(inout) :: C !wavefield in (w,x) domain
    real   , dimension (nx) , intent(in)    :: v ! 0.5*interval velocity
    real                                    :: c1, c2, beta
    complex                                 :: alpha
    complex, dimension(:), allocatable      :: a, b, cc, r, u

    real                                    :: v0, f! reference velocity
    real                                    :: sqa !sqrt(kz)

    allocate(a(nx), b(nx), cc(nx), r(nx), u(nx));


    !! reference velocity
    v0 =minval(v)
!    v0=sum(v)/k%n


    if (mode .eq. 1) then    ! Set coeff for 15 degree eqn
       c1=0.5; c2=0.0;
    elseif (mode .eq. 2) then   ! Set coeff for 45 degree eqn    
       c1=0.5; c2=0.25;   
    else                       ! Set coeff for 60 degree eqn
       c1=0.47824206; c2=0.37636953;
    endif

    do iw = 1, w%n
       
       do ix=1,nx     ! Clear temp array
          a(ix)=0.; b(ix)=0.; cc(ix)=0.; 
          r(ix)=0.; u(ix)=0.;
       enddo

       !!--Phase Shift step
       !!.. forward FFT along x
       call FFT1DX(C(iw,:), +1)
       !Operations in the wavenumber domain
       do ik = 1, k%n
          !---phase shift
          sqa = w%all(iw)**2 / v0**2 - k%all(ik)**2
          if ( sqa > 0 ) then
             C(iw,ik) = C(iw,ik) * exp( i * sign * dz * sqrt(sqa) )
          else
             C(iw,ik) = 0.
          end if
       end do
       !!.. inverse FFT along x
       call FFT1DX(C(iw,:), -1)

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

       !---Fourier FD   

       do ix=2,nx-1
          alpha=-i*sign*c1*v0*dz/(2.*w%all(iw)*dx**2)
!          alpha= i*(-c2*v0**2 + (c1-c2)*v(ix)*v0 - c2*v(ix)**2 )*dz/(2.*w%all(iw)*v(ix)*dx**2);
          beta = c2*( v(ix)**2 + v(ix)*v0 + v0**2)/(w%all(iw)*dx)**2;


          a(ix)=beta-alpha;
          cc(ix)=beta-alpha;
          b(ix)=1-2.*beta+2.*alpha;

          r(ix) = (beta + alpha)*C(iw,ix-1) &
               +(1 - 2.*beta - 2.*alpha)*C(iw,ix) & 
               +(beta + alpha)*C(iw,ix+1);
       enddo

       a(1)=1.;  b(1)=1.;  cc(1)=1.;
       a(nx)=1.;  b(nx)=1.; cc(nx)=1.;

       r(1)=0.;   
       r(nx)=0.;

       ! Call tridiagonal solver
       call ctridag(a,b,cc,r,u,nx)
       

       C(iw,:) = u;


    end do

    deallocate(a,b,cc,r,u);


  end subroutine post3dvz_ffd



  !!---------------------------------------------------------------------------
  !!.. complex tridiagonal solver
  !!---------------------------------------------------------------------------
  !!.. a(n), sub-diagonal elemnets, only elements from 2:n are used
  !!.. b(n), diagonal elemetns
  !!.. c(n), super-diagonal elements, only elements from 1:n-1 are used
  !!
  !!   /  b(1), c(1)                                     \ / u(  1) \     / r(  1) \ 
  !!   |  a(2), b(2), c(2)                               | | u(  2) |     | r(  2) |
  !!   |        a(3), b(3), c(3)                         | | u(  3) |  =  | r(  3) |
  !!   |              ... ...                            | |  ...   |     |  ...   |
  !!   |                          a(n-1), b(n-1), c(n-1) | | u(n-1) |     | r(n-1) |
  !!   \                                  a(  n), b(  n) / \ u(  n) /     \ r(  n) /
  
  subroutine ctridag(a,b,c,r,u,n)
    integer            :: n
    complex            :: a(n), b(n), c(n), r(n), u(n)
    integer            :: j
    integer, parameter :: NMAX = 5000
    complex            :: bet, gam(NMAX)
    bet  = b(1)
    u(1) = r(1)/bet
    do j=2, n
       gam(j) = c(j-1)/bet
       bet    = b(j)-a(j)*gam(j)
       u(j) = (r(j)-a(j)*u(j-1))/bet
    end do
    do j=n-1,1,-1
       u(j) = u(j) - gam(j+1)*u(j+1)
    end do
    return
  end subroutine


!--------------------------------------------------------
! Claudio Version of Finite Differences


subroutine post3dvz_fd2(P, v, mode)
    ! 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,mode
    real                                         :: coef1, coef2,coef3
    real                                        :: vw
    complex                                :: url,urr,ror,rol
    complex, dimension(nx)                        :: U,ra,rb,rc 

   if (mode.eq.1) then !---'finite-difference 15deg'
     coef1=0.5;coef2=0.;coef3=0.
   elseif (mode.eq.2) then !---'finite-difference 45deg'
     coef1=0.5;coef2=0.25;coef3=0.
   else !---'finite-difference 60deg'
     coef1=0.4765;coef2=0.3725;coef3=0.11
   end if
   U=0
   do iw = 1, w%n
      P(iw,:) = P(iw,:)*exp(i*sign*dz*w%all(iw)/v)
      do ix=1,nx
        if (ix==1)  then
          url=0.0
        else
          url=P(iw,ix-1)
        end if
        if (ix==nx) then
          urr=0.0
        else
          urr=P(iw,ix+1)
        end if
        vw=v(ix)/(w%all(iw))
        ror=cmplx(coef3+coef2*vw*vw/(dx*dx),sign*coef1*vw*dz/(2.0*dx*dx)) 
        rol=conjg(ror)
        u(ix)=P(iw,ix) +ror*(url+urr-2.0*P(iw,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)
        P(iw,:)=u
   end do

end subroutine post3dvz_fd2

!--------------------------------------------------------
! Fourier Fd with claudios code

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

  integer                                 :: ik, iw, ix, mode
  complex, dimension (:,:), intent(inout) :: C !wavefield in (w,x) domain
  real   , dimension (nx) , intent(in)    :: v ! 0.5*interval velocity
  real                                         :: coef1, coef2,coef3
  real                                        :: vw, vw1, vw0
  complex                                :: url,urr,ror,rol
  complex, dimension(nx)                        :: U,ra,rb,rc 
  
  real                                    :: v0, vref! reference velocity
  real                                    :: sqa !sqrt(kz)



  !! reference velocity
  v0 =sum(v)/k%n


  do iw = 1, w%n
       

     !!--Phase Shift step
     !!.. forward FFT along x
     call FFT1DX(C(iw,:), +1)
     !Operations in the wavenumber domain
     do ik = 1, k%n
        !---phase shift
        sqa = w%all(iw)**2 / v0**2 - k%all(ik)**2
        if ( sqa > 0 ) then
           C(iw,ik) = C(iw,ik) * exp( i * sign * dz * sqrt(sqa) )
        else
           C(iw,ik) = 0.
        end if
     end do
     !!.. inverse FFT along x
     call FFT1DX(C(iw,:), -1)

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

     !---Fourier FD
     
     if (mode.eq.1) then !---'finite-difference 15deg'
        coef1=0.5;coef2=0.;coef3=0.
     elseif (mode.eq.2) then !---'finite-difference 45deg'
        coef1=0.5;coef2=0.25;coef3=0.
     else !---'finite-difference 60deg'
        coef1=0.4765;coef2=0.3725;coef3=0.11
     end if

     U=0

      
     do ix=1,nx
        if (ix==1)  then
           url=0.0
        else
           url=C(iw,ix-1)
        end if
        if (ix==nx) then
           urr=0.0
        else
           urr=C(iw,ix+1)
        end if
        !          vw=v(ix)/(w%all(iw))
        vw1=v(ix)/w%all(iw)
        vw0=v0/w%all(iw)*0.02
        ror=cmplx(coef3+coef2*(vw1**2+vw0*vw1+vw0**2)/(dx*dx),sign*coef1*vw0*dz/(2.0*dx*dx)) 
        rol=conjg(ror)
        u(ix)=C(iw,ix) +ror*(url+urr-2.0*C(iw,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)
     C(iw,:)=u
      
  end do
    
end subroutine post3dvz_ffd2

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