module sogrot_mod2

  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
  public :: sogrot_init, sogrot_exe, sog_fft, sog_ifft, rotsou_geodip, rotrec_geodip
  
  real,    parameter 			:: pi=3.1415926535898
  complex, parameter 			:: i = cmplx(0.,1.)
  integer            			:: nx, nh, nz
  integer	     			:: npadx, npadz, npadh
  real               			:: dx, dh, dz
  real               			:: ox, oh, oz
  type(fft_type)     			:: fft_1d_idx, fft_1d_idz, fft_1d_idh, fft_2d_idxh, fft_2d_idzh
  type(axis)         			:: kh, kx, 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, h) result(stat)

    type(axis), intent(in) :: x, z, h
    integer                :: ih, ikx, ikz, stat
    type(axis)             :: ks

    nx = x%n;    dx = x%d
    nz = z%n;    dz = z%d
    nh = h%n;    dh = h%d

    stat=1
!!!
    kh%n = nh+npadh
    kh%d = 1./(kh%n*dh)
    kh%o = 1./(2*dh)-kh%d ! no zero frequency component
    allocate(kh%all(kh%n)); kh%all=0.
    do ih = 1, kh%n
      kh%all(ih) = mod(kh%o+(ih-1)*kh%d,1./dh)-kh%o
    end do
    kh%all = 2*pi * kh%all
!!!
    kx%n = nx+npadx
    kx%d = 1./(kx%n*dx)
    kx%o = 1./(2*dx)-kx%d !k does not increase uniformly!
    allocate(kx%all(kx%n)); kx%all=0.
    do ikx=1, kx%n
       kx%all(ikx) = mod(kx%o+(ikx-1)*kx%d,1./dx)-kx%o
    end do
    kx%all = 2*pi * kx%all
!!!
    kz%n = nz+npadz
    kz%d = 1./(kz%n*dz)
    kz%o = 1./(2*dz)-kz%d !kz does not increase uniformly!
    allocate(kz%all(kz%n)); kz%all=0.
    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
!!!
    fft_1d_idx = init_fft_1d(kx%n)
    fft_1d_idz = init_fft_1d(kz%n)
    fft_2d_idxh = init_fft_2d(kh%n,kx%n)
    fft_2d_idzh = init_fft_2d(kz%n,kh%n)

!    write(0,*) kz%all ! n,kz%o,kz%d
!    write(0,*) kh%all ! n,kh%o,kh%d
!    write(0,*) kx%all ! n,kx%o,kx%d

stat = 0

  end function fftvar_init
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 subroutine sogrot_init(npadz_in,npadx_in,npadh_in)

 integer			:: npadz_in,npadx_in,npadh_in

 npadz=npadz_in; npadx=npadx_in; npadh=npadh_in

 end subroutine sogrot_init 
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function sogrot_exe(cig,souwav,push) result(cig_out)

 logical				:: souwav,push
 real, intent(inout)			:: cig(:,:,:)
 complex,allocatable			:: Ccig(:,:,:),Rcig(:,:,:)
 real,allocatable			:: cig_out(:,:,:)

 allocate(Ccig(kz%n,kh%n,kx%n),Rcig(kz%n,kh%n,kx%n))
 allocate(cig_out(nz,nh,nx))

 write(0,*) size(cig)
 Ccig = sog_fft(cig) 
 if (souwav) then
    Rcig = rotsou_geodip(Ccig)
 else
    Rcig = rotrec_geodip(Ccig)
 end if
 cig_out  = sog_ifft(Rcig)
 write(0,*) size(cig_out)
 
 end function sogrot_exe
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function sog_fft(cig) result(Ccig)

 complex, allocatable			:: Ccig(:,:,:)
 real, intent(in)			:: cig(:,:,:)
 integer				:: ix, iz, ih

 allocate(Ccig(kz%n,kh%n,kx%n))
 Ccig=cmplx(0,0)
 Ccig(1:nz,1:nh,1:nx)=cmplx(cig,0)
 do ix=1,kx%n
    do ih=1,kh%n
       call fft_1d_data(fft_1d_idz,  1,  kz%n,Ccig(:,ih,ix))
   end do
 end do
 
 do iz=1,kz%n
    call fft_2d_data(fft_2d_idxh,  1,  kh%n,kx%n,Ccig(iz,:,:))
 end do

 end function sog_fft
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function sog_ifft(Ccig) result(cig)

 complex, intent(in)			:: Ccig(:,:,:)
 real, allocatable			:: cig(:,:,:)
 complex, allocatable			:: temp(:,:,:)
 integer				:: ix, iz, ih

 allocate(cig(nz,nh,nx))
 allocate(temp(kz%n,kh%n,kx%n))

 cig=0

 temp=Ccig

 do iz=1,kz%n
    call fft_2d_data(fft_2d_idxh,  -1,  kh%n, kx%n,Ccig(iz,:,:))
 end do

 do ix=1,kx%n
    do ih=1,kh%n
       call fft_1d_data(fft_1d_idz,  -1,  kz%n,Ccig(:,ih,ix))
   end do
 end do

 cig=real(Ccig(1:nz,1:nh,1:nx))

 end function sog_ifft
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function rotsou_geodip(Ccig) result(Scig)

 complex, intent(in)			:: Ccig(:,:,:)
 complex, allocatable			:: Scig(:,:,:)
 integer				:: ix, iz, ih, stat, ikhrot
 real					:: tgama,talfa,tteta_s

 allocate(Scig(kz%n,kh%n,kx%n));Scig=cmplx(0,0)

 do ix=1,kx%n
    do iz=1,kz%n
       if(kz%all(iz).eq.0.) then
          if(kx%all(ix).lt.0.) then
             talfa=99999999999.
          else
             talfa=-99999999999.
          end if
       else
          talfa=-kx%all(ix)/kz%all(iz)
       end if
       do ih=1,kh%n
          if(kz%all(iz).eq.0.) then
             if(kh%all(ih).lt.0.) then
                tgama=99999999999.
             else
                tgama=-99999999999.
             end if
          else
             tgama=-kh%all(ih)/kz%all(iz)
          end if

	  if (1-talfa*tgama.eq.0.) then
	     tteta_s=99999999999.
	  else
             tteta_s=(talfa+tgama)/(1-talfa*tgama)
	  end if

!          write(2,*) " tgama = ",tgama," talfa = ",talfa," tteta_s = ",tteta_s

	  ikhrot=int(kz%all(iz)*tteta_s/kh%d-kh%o)+1
	   
!          write(2,*) " iz = ",iz," ih = ",ih," ix = ",ix," ikhrot = ",ikhrot
	  if (ikhrot.ge.1 .and. ikhrot.le.kh%n) then
  	     Scig(iz,ikhrot,ix)=Scig(iz,ikhrot,ix)+Ccig(iz,ih,ix)
!  	     Scig(iz,ih,ix)=Scig(iz,ih,ix)+Ccig(iz,ikhrot,ix)
	  end if
       end do
    end do
 end do

 end function rotsou_geodip
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function rotrec_geodip(Ccig) result(Gcig)

 complex, intent(in)			:: Ccig(:,:,:)
 complex, allocatable			:: Gcig(:,:,:)
 integer				:: ix, iz, ih, stat, ikhrot
 real					:: tgama,talfa,tteta_g

 allocate(Gcig(kz%n,kh%n,kx%n));Gcig=cmplx(0,0)

 do iz=1,kz%n
    do ih=1,kh%n
       if(kz%all(iz).eq.0.) then
	  if(kh%all(ih).lt.0.) then
             tgama=99999999999.
	  else
             tgama=-99999999999.
	  end if
       else
	  tgama=-kh%all(ih)/kz%all(iz)
       end if

       do ix=1,kx%n
          if(kz%all(iz).eq.0.) then
	     if(kx%all(ih).lt.0.) then
    	        talfa=99999999999.
	     else
                talfa=-99999999999.
	     end if
          else
	     talfa=-kx%all(ix)/kz%all(iz)
	  end if

	  if (1+talfa*tgama.eq.0.) then
	     tteta_g=99999999999.
	  else
             tteta_g=(talfa-tgama)/(1+talfa*tgama)
	  end if

!          write(0,*) " tgama = ",tgama," talfa = ",talfa," tteta_s = ",tteta_s," tteta_g = ",tteta_g
	  ikhrot=int(-kz%all(iz)*tteta_g/kh%d-kh%o)+1
!          write(0,*) " iz = ",iz," ih = ",ih," ix = ",ix," ikhrot = ",ikhrot
	  if (ikhrot.ge.1 .and. ikhrot.le.kh%n) then
  	     Gcig(iz,ikhrot,ix)=Gcig(iz,ikhrot,ix)+Ccig(iz,ih,ix)
	  end if
       end do
    end do
 end do

 end function rotrec_geodip
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
end module sogrot_mod2
