module sogrot_mod6

  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_I_shift, sogrot_F_shift
  public :: sogrot_init, sogrot_exe, sog_fft, sog_ifft, rotsou_geodip, rotrec_geodip, check_sum, rotsou_chk
  
  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; oh = h%o

    stat=1
!!!
!    kh%n = nh+npadh
!    kh%d = 2*pi / (kh%n*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)
    kh%n = nh+npadh
    kh%d = 2*pi / (kh%n*dh)
    allocate(kh%all(kh%n)); kh%all=0.
    do ikh = 1,kh%n/2
       kh%all(ikh       ) = (        ikh-1)
       kh%all(ikh+kh%n/2) = (-kh%n/2+ikh-1)
    end do
    kh%all = 2*pi* kh%all / (kh%n*dx)
    kh%o = kh%all(1)
!!!
    kx%n = nx+npadx
    kx%d = 2*pi / (kx%n*dx)
    allocate(kx%all(kx%n)); kx%all=0.
    do ikx = 1,kx%n/2
       kx%all(ikx       ) = (        ikx-1)
       kx%all(ikx+kx%n/2) = (-kx%n/2+ikx-1)
    end do
    kx%all = 2*pi* kx%all / (kx%n*dx)
    kx%o = kx%all(1)
!!!
    kz%n = nz+npadz
    kz%d = 2*pi / (kz%n*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)
    kz%o = kz%all(1)
!!!
    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_I_shift(f)

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

 do ih=1,kh%n
    f(ih) = f(ih) * cexp(cmplx(0,kh%all(ih)*oh))
 end do

 end subroutine sogrot_I_shift
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 subroutine sogrot_F_shift(f)

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

 do ih=1,kh%n
    f(ih) = f(ih) * cexp(cmplx(0,-kh%all(ih)*oh))
 end do

 end subroutine sogrot_F_shift

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function check_sum(a) result(b)

 real					:: a(:,:,:)
 real					:: b
 integer				:: iz,ih,ix

 b=0
 do ix=1,size(a,3); do ih=1,size(a,2); do iz=1,size(a,1)
 b=b+abs(a(iz,ih,ix))
 end do; end do; end do
 
 end function check_sum
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 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
 real					:: b

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

 Rcig=0.;cig_out=0.

 if (souwav) then
    Rcig = rotsou_geodip(Ccig)
 else
    Rcig = rotrec_geodip(Ccig)
 end if
 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,kh%n,kx%n))
 cig=0.

 do ix=1,kx%n
    do ih=1,kh%n
!        do iz=kz%n/2+2,kz%n
!          Ccig(iz,ih,ix)=conjg(Ccig(kz%n-iz+2,ih,ix))
!        end do
       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,:,:))

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

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

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

 do ix=1,kx%n
    do iz=1,kz%n
       if(iz.eq.1) 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

       do ih=1,kh%n
          if(iz.ne.1) then
             tteta_s=-kh%all(ih)/kz%all(iz)
	     if(1+tteta_s*talfa > tiny(0.)) then
	        tgama = -(talfa-tteta_s)/(1+tteta_s*talfa)
	        khrot=-kz%all(iz)*tgama
	     else
                if(kz%all(iz).lt.0.) then
		   khrot=huge(1.)!kh%all(ih)
                else if(kz%all(iz).gt.0.) then
                   khrot=-huge(1.)!kh%all(ih)
                else
		   khrot=0.
                end if
	     end if
	  else
	     khrot=0.
	  end if	     

          if (khrot.ge.kh%all(1).and.khrot.le.kh%all(kh%n/2)) then
             ikhrot=int(khrot/kh%d)+1
          else
             ikhrot=kh%n+int(khrot/kh%d)+1
          end if
             
	  fwrot=(khrot)/kh%d-ikhrot-1
	  gwrot=1.-fwrot

     	  if (ikhrot.gt.1 .and. ikhrot.le.kh%n/2-1) then
	     k=k+1
	     Scig(iz,ih,ix)=gwrot*Ccig(iz,ikhrot,ix)+fwrot*Ccig(iz,ikhrot+1,ix)
	  else if (ikhrot.gt.kh%n/2 .and. ikhrot.le.kh%n-1) then
	     k=k+1
	     Scig(iz,ih,ix)=gwrot*Ccig(iz,ikhrot,ix)+fwrot*Ccig(iz,ikhrot+1,ix)
	  else if (ikhrot.eq.1 .or. ikhrot.eq.kh%n) then
	     k=k+1
	     Scig(iz,ih,ix)=gwrot*Ccig(iz,1,ix)+fwrot*Ccig(iz,kh%n,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(:,:,:),tempk(:),tempr(:)
 integer				:: ix, iz, ih, stat, ikhrot,k
 real					:: tgama,talfa,tteta_g,gwrot,fwrot,khrot

 allocate(Gcig(kz%n,kh%n,kx%n),tempk(kh%n),tempr(kh%n));Gcig=0.
 k=0
 do ix=1,kx%n
    do iz=1,kz%n/2
       if(iz.ne.1) then
	  talfa=-kx%all(ix)/kz%all(iz)
       else
          if(kx%all(ix).le.0.) then
             talfa=-huge(1.)
          else
             talfa=huge(1.)
          end if
       end if

       do ih=1,kh%n
          if(iz.ne.1) then
             tteta_g=-kh%all(ih)/kz%all(iz)
	     if(1-tteta_g*talfa > tiny(0.)) then
	        tgama = (talfa+tteta_g)/(1-tteta_g*talfa)
	        khrot=-kz%all(iz)*tgama
	     else
		khrot=kh%all(ih)
	     end if
	  else
	     khrot=kh%all(ih)
	  end if	     

          if (khrot.ge.kh%all(1).and.khrot.le.kh%all(kh%n/2)) then
             ikhrot=int(khrot/kh%d)+1
          else
             ikhrot=kh%n+int(khrot/kh%d)+1
          end if
             
	  fwrot=(khrot)/kh%d-ikhrot-1
	  gwrot=1.-fwrot

     	  if (ikhrot.ge.1 .and. ikhrot.le.kh%n/2-1) then
	     k=k+1
	     Gcig(iz,ih,ix)=gwrot*Ccig(iz,ikhrot,ix)+fwrot*Ccig(iz,ikhrot+1,ix)
	  else if (ikhrot.ge.kh%n/2 .and. ikhrot.le.kh%n-1) then
	     k=k+1
	     Gcig(iz,ih,ix)=gwrot*Ccig(iz,ikhrot,ix)+fwrot*Ccig(iz,ikhrot+1,ix)
	  end if
       end do

    end do
 end do
 write(0,*) k

 write(0,*) "passei rotrec_geodip"

 end function rotrec_geodip
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function rotsou_chk(Ccig) result(Scig)

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

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

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

       do ih=1,kh%n/2+1
	  tempk(ih+kh%n/2-1)=Ccig(iz,ih,ix)
       end do
       do ih=kh%n/2+2,kh%n
	  tempk(ih-kh%n/2-1)=Ccig(iz,ih,ix)
       end do
       do ih=1,kh%n
          tempk(ih)=tempk(ih)*cexp(cmplx(0,kh%all(ih)*oh))
       end do

!       do ih=1,kh%n
!          if(iz.ne.1) then
!             tteta_s=-kh%all(ih)/kz%all(iz)
!	     if(1+tteta_s*talfa > tiny(0.)) then
!	        tgama = (talfa-tteta_s)/(1+tteta_s*talfa)
!	        khrot=-kz%all(iz)*tgama
!	     else
!		khrot=kh%all(ih)
!	     end if
!	  else
!	     khrot=kh%all(ih)
!	  end if	     

!          ikhrot=int(khrot/kh%d)+int(kh%n/2)+1
!	  fwrot=(khrot)/kh%d-ikhrot-1
!	  gwrot=1.-fwrot

!     	  if (ikhrot.ge.1 .and. ikhrot.lt.kh%n) then
!	     k=k+1
!	     tempr(ih)=gwrot*tempk(ikhrot)+fwrot*tempk(ikhrot+1)!+tempr(ih)
!	  else if (ikhrot.eq.kh%n) then
!	     k=k+1
!	     tempr(ih)=gwrot*tempk(ikhrot)+fwrot*tempk(1)!+tempr(ih)
!	  else if  (ikhrot.eq.0) then
!	     k=k+1
!	     tempr(ih)=gwrot*tempk(kh%n)+fwrot*tempk(1)!+tempr(ih)
!	  end if
!       end do

!       do ih=1,kh%n
!          tempr(ih)=tempr(ih)*cexp(cmplx(0,-kh%all(ih)*oh))
!       end do

!       do ih=1,kh%n/2+1
!	  Scig(iz,ih,ix)=tempr(ih+kh%n/2-1)
!       end do
!       do ih=kh%n/2+2,kh%n
!	  Scig(iz,ih,ix)=tempr(ih-kh%n/2-1)
!       end do
       do ih=1,kh%n
!	  Scig(iz,ih,ix)=tempr(ih)
	  Scig(iz,ih,ix)=tempk(ih)
       end do
    end do
 end do
 write(0,*) k
! Scig=Ccig
 write(0,*) "passei rotsou_chk"
 end function rotsou_chk
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
end module sogrot_mod6
