module sogrot_mod4

  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, sogrot_shift_init, sogrot_shift_close,sogrot_I_shift, sogrot_F_shift
  public :: sogrot_init, sogrot_exe, sog_fft, sog_ifft, rotsou_geodip, rotrec_geodip
  
!  complex, dimension(:,:,:),allocatable:: sftF
!  complex, dimension(:,:,:),allocatable:: sftI
  complex, dimension(:),allocatable 	:: sftF
  complex, dimension(:),allocatable 	:: sftI
  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,fft_3d_idzhx
  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                :: ikh, ikx, ikz, stat

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

    stat=1
!!!
    kh%n = nh
    kh%d = 2*pi / (nh*dh)
    allocate(kh%all(kh%n)); kh%all=0.
    do ikh = 1,kh%n
       kh%all(ikh) = (-kh%n/2+ikh-1)
    end do
    kh%all = 2*pi* kh%all / (kh%n*dh)
    kh%o = kh%all(1)
!!!
    kx%n = nx
    kx%d = 2*pi / (nx*dx)
    allocate(kx%all(kx%n)); kx%all=0.
    do ikx = 1,kx%n
       kx%all(ikx) = (-kx%n/2+ikx-1)
    end do
    kx%all = 2*pi* kx%all / (kx%n*dx)
    kx%o = kx%all(1)
!!!
    kz%n = nz
    kz%d = 2*pi / (nz*dz)
    allocate(kz%all(kz%n)); kz%all=0.
    do ikz = 1,kz%n/2
       kz%all(ikz       ) = (        ikz-1)
       kz%all(ikz+kz%n/2) = (-kz%n/2+ikz-1)
    end do
    kz%all = 2*pi* kz%all / (kz%n*dz)
!!!
    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)
    fft_3d_idzhx = init_fft_3d(kz%n,kh%n,kx%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 
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!  subroutine sogrot_shift_init(kz,kh,kx)
  subroutine sogrot_shift_init(kh)

!    type(axis), intent(in):: kz,kh,kx
    type(axis), intent(in):: kh
    integer             :: i1,i2,i3
    real                :: pi, argF, argI
    integer             :: stat

!    allocate(sftI(kz%n,kh%n,kx%n))
!    allocate(sftF(kz%n,kh%n,kx%n))

    allocate(sftI(kh%n))
    allocate(sftF(kh%n))

    sftF=cmplx(1.,0.)
    sftI=cmplx(1.,0.)

!    if(kz%o/=0. .or. kh%o/=0. .or. kx%o/=0 ) then
    if(kh%o/=0.) then
       pi = acos(-1.)
!       do i1=1,kz%n
          do i2=1,kh%n
!             do i3=1,kx%n

!                argI = 2*pi &
!                * ( (i1-1.)/kz%n * kz%o/kz%d &
!                +   (i2-1.)/kh%n * kh%o/kh%d &
!                +   (i3-1.)/kx%n * kx%o/kx%d )

                argI = 2*pi*(i2-1.)/kh%n * kh%o/kh%d 

                argF = 2*pi*(i2-1.)/kh%n * kh%o/kh%d

!                argF = 2*pi &
!                * ( (i1-1.)/kz%n * kz%o/kz%d &
!                +   (i2-1.)/kh%n * kh%o/kh%d &
!                +   (i3-1.)/kx%n * kx%o/kx%d )

!                sftI(i1,i2,i3) = cmplx(cos(argI),-(+1)*sin(argI))
!                sftF(i1,i2,i3) = cmplx(cos(argF),-(-1)*sin(argF))

                sftI(i2) = cmplx(cos(argI),-(+1)*sin(argI))
                sftF(i2) = cmplx(cos(argF),-(-1)*sin(argF))

!             end do
!          end do
       end do

    end if
 write(0,*) "passei sogrot_shift_init"
  end subroutine sogrot_shift_init
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine sogrot_I_shift(f)

  complex, dimension(:,:,:)			:: f
  integer					:: ih

 do ih=1,kh%n
    f(:,ih,:) = f(:,ih,:) * sftI(ih)
 end do

 write(0,*) "passei sogrot_I_shift"
  end subroutine sogrot_I_shift
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine sogrot_F_shift(f)

  complex, dimension(:,:,:)			:: f
  integer					:: ih

 do ih=1,kh%n
    f(:,ih,:) = f(:,ih,:) * sftF(ih)
 end do

 write(0,*) "passei sogrot_F_shift"
  end subroutine sogrot_F_shift
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  subroutine sogrot_shift_close()

    deallocate(sftI,sftF)

 write(0,*) "passei sogrot_shift_close"
  end subroutine sogrot_shift_close
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function sogrot_exe(Ccig,souwav) result(cig_out)

 logical				:: souwav
 complex, intent(in)			:: Ccig(:,:,:)
 real, allocatable			:: cig_out(:,:,:)
 complex,allocatable			:: Rcig(:,:,:)
 integer				:: ix, iz, ih

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

 Rcig=0.;cig_out=0.

 if (souwav) then
    Rcig = rotsou_geodip(Ccig)
 else
    Rcig = rotrec_geodip(Ccig)
 end if
 call sogrot_I_shift(Rcig)
 cig_out = sog_ifft(Rcig)

 deallocate (Rcig)

 write(0,*) "passei sogrot_exe"
 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))

 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

 write(0,*) "passei sog_fft"
 end function sog_fft
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function sog_ifft(Ccig) result(cig) 

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

 allocate(cig(nz,nh,nx))
 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

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

 write(0,*) "passei sog_ifft"
 end function sog_ifft
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function rotsou_geodip(Ccig) result(Scig)

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

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

! write(2,*) "   kx     ","   kh     ","   kz     ","  talfa   ","  tgama   ","  tetas   ","   ikhrot "

 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(1.)
             else
                talfa=huge(1.)
             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(1.)
             else
                tgama=-huge(1.)
             end if
          else
             tgama=-kh%all(ih)/kz%all(iz)
          end if
!	  talfa=0.
	  if (1.-talfa*tgama.eq.0.) then
	     if(talfa+tgama.le.0.) then
	       tteta_s=-huge(1.)
	     else
	       tteta_s=huge(1.)
	     end if
	  else
             tteta_s=(talfa+tgama)/(1.-talfa*tgama)
	  end if
	
!	  ikhrot=int((-kz%all(iz)*tteta_s-kh%o)/kh%d)+1
	  ikhrot=int((-kz%all(iz)*tteta_s)/kh%d)+1

!	  fwrot=(-kz%all(iz)*tteta_s-kh%o)/kh%d-ikhrot-1
	  fwrot=(-kz%all(iz)*tteta_s)/kh%d-ikhrot-1
	  gwrot=1.-fwrot

  	  if (ikhrot.ge.1 .and. ikhrot.lt.kh%n) then
	     k=k+1
  	     Scig(iz,ikhrot,ix)=Scig(iz,ikhrot,ix)+gwrot*Ccig(iz,ih,ix)
  	     Scig(iz,ikhrot+1,ix)=Scig(iz,ikhrot+1,ix)+fwrot*Ccig(iz,ih,ix)
	  end if
       end do
    end do
 end do
 write(0,*) k

 write(0,*) "passei rotsou_geodip"
 end function rotsou_geodip
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function rotrec_geodip(Ccig) result(Gcig)

 complex				:: Ccig(:,:,:)
 complex, allocatable			:: Gcig(:,:,:)
 integer				:: ix, iz, ih, stat, ikhrot,k
 real					:: tgama,talfa,tteta_g,gwrot,fwrot

 allocate(Gcig(kz%n,kh%n,kx%n));Gcig=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(1.)
             else
                talfa=huge(1.)
             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(1.)
             else
                tgama=-huge(1.)
             end if
          else
             tgama=-kh%all(ih)/kz%all(iz)
          end if
!	  talfa=0.
	  if (1.+talfa*tgama.eq.0.) then
	     if(talfa+tgama.le.0.) then
	       tteta_g=-huge(1.)
	     else
	       tteta_g=huge(1.)
	     end if
	  else
             tteta_g=(talfa-tgama)/(1.+talfa*tgama)
	  end if

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

!	  fwrot=(kz%all(iz)*tteta_g-kh%o)/kh%d-float(ikhrot)
	  fwrot=(kz%all(iz)*tteta_g)/kh%d-float(ikhrot)
	  gwrot=1.-fwrot

  	  if (ikhrot.ge.1 .and. ikhrot.lt.kh%n) then
	     k=k+1
  	     Gcig(iz,ikhrot,ix)=Gcig(iz,ikhrot,ix)+gwrot*Ccig(iz,ih,ix)
  	     Gcig(iz,ikhrot+1,ix)=Gcig(iz,ikhrot+1,ix)+fwrot*Ccig(iz,ih,ix)
	  end if
       end do
    end do
 end do
 write(0,*) k

 write(0,*) "passei rotrec_geodip"
 end function rotrec_geodip
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
end module sogrot_mod4
