module sogrot_mod

!  use fftest_type_mod
  use fft_data_mod
  implicit none

  type axis
     integer 					:: n
     real    					:: o, d, v ! v=(i-1)*d+o
     real, allocatable				:: all(:)
  end type axis

  public :: axis
  private
  public :: fftvar_init,sogrot
  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
  real               :: ox, oz, ot
  type(fft_type)     :: fft_1d_id, fft_2d_id, fft_3d_id, rft_1d_id
  type(axis)         :: w, k, kz

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

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

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

    sign=-1 ! if(dowcon_as_shot) sign = +1
    nx = x%n;    dx = x%d
    nt = t%n;    dt = 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
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    fft_1d_id = init_fft_1d(nz)
    fft_2d_id = init_fft_2d(nt,nz) 

! write(0,*) fft_1d_id%ierr,fft_1d_id%ndim,fft_1d_id%plan_forw,fft_1d_id%plan_back,fft_1d_id%scale

    kz%n = nz
    allocate(kz%all(kz%n)); kz%all=0.

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

    kz%all = 2*pi * kz%all

  end function fftvar_init
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  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,nc

    nc=size(C,2)
    all_freq (1)=0
    do ix = 1, nc
       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 sogrot(u)

 complex, allocatable			:: C(:,:)
 real, intent(inout)			:: u(:,:)
 integer				:: iw

 allocate(C(nt,kz%n))

 C = cmplx(u,0)
 call fft_2d_data(fft_2d_id,   1,  nt, kz%n,C)
 C(int(nt/2)+1:nt,int(kz%n/2)+1:kz%n)=0.
 C(1:int(nt/2)+1,1:int(kz%n/2)+1)=0.
 call fft_2d_data(fft_2d_id,  -1,  nt, kz%n,C)
 u=real(C)

 end subroutine sogrot

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
end module sogrot_mod
