module sogrot_mod3_fft

  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
  public :: kh,kx,kz
  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 = 2*pi/(kh%n*dh)
    kh%o = -(kh%n/2 -1)*kh%d
    allocate(kh%all(kh%n)); kh%all=0.
    do ih = 1, kh%n
      kh%all(ih) = kh%o+(ih-1)*kh%d
    end do
!!!
    kx%n = nx+npadx
    kx%d = 2*pi/(kx%n*dx)
    kx%o = -(kx%n/2 -1)*kx%d
    allocate(kx%all(kx%n)); kx%all=0.
    do ikx = 1, kx%n
      kx%all(ikx) = kx%o+(ikx-1)*kx%d
    end do
!!!
    kz%n = nz+npadz
    kz%d = 2*pi/(kz%n*dz)
    kz%o = -(kz%n/2 -1)*kz%d
    allocate(kz%all(kz%n)); kz%all=0.
    do ikz = 1, kz%n
      kz%all(ikz) = kz%o+(ikz-1)*kz%d
    end do
!!!
    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%n,kz%o,kz%d
!    write(0,*) kh%n,kh%o,kh%d
!    write(0,*) kx%n,kx%o,kx%d

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

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(in)			:: cig(:,:,:)
 real, allocatable			:: cig_out(:,:,:)
 complex,allocatable			:: Ccig(:,:,:),Rcig(:,:,:)
 integer				:: ix, iz, ih

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

 Ccig=0.;Rcig=0.;cig_out=0.

 Ccig = sog_fft(cig) 
 if (souwav) then
    Rcig = rotsou_geodip(Ccig,push)
 else
    Rcig = rotrec_geodip(Ccig,push)
 end if
 cig_out = sog_ifft(Rcig)

 deallocate (Ccig,Rcig)
 end function sogrot_exe
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function sog_fft(cig) result(Ccig)

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

 allocate(Ccig(kz%n,kh%n,kx%n))

 do ix=1,nx
    do ih=1,nh
       do iz=1,nz
          cig(iz,ih,ix)=cig(iz,ih,ix)*(-1)**(ix+ih+iz)
       end do
    end do
 end do
 Ccig=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) !result(cig) 

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

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

 temp=Ccig
 cig=0.
 do ix=1,kx%n
    do ih=1,kh%n
       do iz=1,kz%n
          if (iz.lt.int(kz%n/2)+1.and.ih.lt.int(kh%n/2)+1.and.ix.lt.int(kx%n/2)+1) then
             temp(iz+int(kz%n/2),ih+int(kh%n/2),ix+int(kx%n/2))=Ccig(iz,ih,ix)
          else if (iz.lt.int(kz%n/2)+1.and.ih.lt.int(kh%n/2)+1.and.ix.ge.int(kx%n/2)+1) then
             temp(iz+int(kz%n/2),ih+int(kh%n/2),ix-int(kx%n/2))=Ccig(iz,ih,ix)
          else if (iz.lt.int(kz%n/2)+1.and.ih.ge.int(kh%n/2)+1.and.ix.ge.int(kx%n/2)+1) then
             temp(iz+int(kz%n/2),ih-int(kh%n/2),ix-int(kx%n/2))=Ccig(iz,ih,ix)
          else if (iz.lt.int(kz%n/2)+1.and.ih.ge.int(kh%n/2)+1.and.ix.lt.int(kx%n/2)+1) then
             temp(iz+int(kz%n/2),ih-int(kh%n/2),ix+int(kx%n/2))=Ccig(iz,ih,ix)
          else if (iz.ge.int(kz%n/2)+1.and.ih.ge.int(kh%n/2)+1.and.ix.ge.int(kx%n/2)+1) then
             temp(iz-int(kz%n/2),ih-int(kh%n/2),ix-int(kx%n/2))=Ccig(iz,ih,ix)
          else if (iz.ge.int(kz%n/2)+1.and.ih.ge.int(kh%n/2)+1.and.ix.lt.int(kx%n/2)+1) then
             temp(iz-int(kz%n/2),ih-int(kh%n/2),ix+int(kx%n/2))=Ccig(iz,ih,ix)
          else if (iz.ge.int(kz%n/2)+1.and.ih.lt.int(kh%n/2)+1.and.ix.lt.int(kx%n/2)+1) then
             temp(iz-int(kz%n/2),ih+int(kh%n/2),ix+int(kx%n/2))=Ccig(iz,ih,ix)
          else if (iz.ge.int(kz%n/2)+1.and.ih.lt.int(kh%n/2)+1.and.ix.ge.int(kx%n/2)+1) then
             temp(iz-int(kz%n/2),ih+int(kh%n/2),ix-int(kx%n/2))=Ccig(iz,ih,ix)
          end if
       end do
    end do
 end do
 Ccig=-temp

 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

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

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

 logical 				:: push
 complex				:: Ccig(:,:,:)
 complex, allocatable			:: Scig(:,:,:)
 integer				:: ix, iz, ih, stat, ikhrot,k
 real					:: tgama,talfa,tteta_s

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

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

          if(kz%all(iz).eq.0.) then
             if(kh%all(ih).le.0.) then
                tgama=huge(float(ix))
             else
                tgama=-huge(float(ix))
             end if
          else
             tgama=-kh%all(ih)/kz%all(iz)
          end if
!	  talfa=0.
	  if (1-talfa*tgama.eq.0.) then
	     tteta_s=huge(float(ix))
	  else
             tteta_s=(talfa+tgama)/(1-talfa*tgama)
	  end if

	  ikhrot=int((-kz%all(iz)*tteta_s-kh%o)/kh%d)+1

  	  if (ikhrot.ge.1 .and. ikhrot.le.kh%n) then
	     k=k+1
  	     Scig(iz,ikhrot,ix)=Scig(iz,ikhrot,ix)+Ccig(iz,ih,ix)
	  end if
       end do
    end do
 end do
 write(0,*) k
 end function rotsou_geodip
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function rotrec_geodip(Ccig,push) result(Gcig)

 logical 				:: push
 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=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

	  ikhrot=int(-kz%all(iz)*tteta_g/kh%d-kh%o)+1

	  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_mod3_fft
