module sogrot_mod3z

!  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, check_sum
  
  real,    parameter 			:: pi=3.1415926535898
  complex, parameter 			:: i = cmplx(0.,1.)
  complex,pointer  			:: tfield(:,:,:)
  integer            			:: nx, nh, nz,nzs(3)
  integer	     			:: npadx, npadz, npadh
  real               			:: dx, dh, dz
  real               			:: ox, oh, oz
  type(axis)         			:: kh, kx, kz

  integer, parameter   			::  FFTW_R2HC=0,FTW_HC2R=1,FFTW_DHT=2,FFTW_REDFT00=3,FFTW_REDFT01=4
  integer, parameter   			::  FFTW_REDFT10=5,FFTW_REDFT11=6,FFTW_RODFT00=7,FFTW_RODFT01=8,FFTW_RODFT10=9
  integer, parameter   			::  FFTW_RODFT11=10,FFTW_FORWARD=-1,FFTW_BACKWARD=+1,FFTW_DESTROY_INPUT=1
  integer, parameter   			::  FFTW_UNALIGNED=2,FFTW_CONSERVE_MEMORY=4,FFTW_EXHAUSTIVE=8,FTW_PRESERVE_INPUT=16
  integer, parameter   			::  FFTW_PATIENT=32,FFTW_ESTIMATE=64,FFTW_ESTIMATE_PATIENT=128
  integer, parameter   			::  FFTW_BELIEVE_PCOST=256,FFTW_DFT_R2HC_ICKY=512,FFTW_NONTHREADED_ICKY=1024
  integer, parameter   			::  FFTW_NO_BUFFERING=2048,FFTW_NO_INDIRECT_OP=4096,FFTW_ALLOW_LARGE_GENERIC=8192
  integer, parameter   			::  FFTW_NO_RANK_SPLITS=16384,FFTW_NO_VRANK_SPLITS=32768,FFTW_NO_VRECURSE=65536
 
  integer*8, private	   		::  planF,planI
  
  ! 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, nh2, nx2, nz2

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

    stat=1

!!!
    kh%n = nh+npadh
    allocate(kh%all(kh%n)); kh%all=0.
    nh2=kh%n/2
    do ikh = 1, nh2
      kh%all(ikh)     = ikh-1
      kh%all(ikh+nh2) = -nh2+ikh-1
    end do
    kh%all = 2*pi*kh%all/(kh%n*dh)
    kh%d = 2*pi/(kh%n*dh)
!!!
    kx%n = nx+npadx
    allocate(kx%all(kx%n)); kx%all=0.
    nx2=kx%n/2
    do ikx = 1, nx2
      kx%all(ikx)     = ikx-1
      kx%all(ikx+nx2) = -nx2+ikx-1
    end do
    kx%all = 2*pi*kx%all/(kx%n*dx)
    kx%d = 2*pi/(kx%n*dx)
!!!
    kz%n = nz+npadz
    allocate(kz%all(kz%n)); kz%all=0.
    nz2=kz%n/2
    do ikz = 1, nz2
      kz%all(ikz)     = ikz-1
      kz%all(ikz+nz2) = -nz2+ikz-1
    end do
    kz%all = 2*pi*kz%all/(kz%n*dz)
    kz%d = 2*pi/(kz%n*dz)
!!!
    nzs=(/kz%n,kh%n,kx%n/)
    allocate(tfield(nzs(1),nzs(2),nzs(3)))
    call sfftw_plan_dft(planF,3,nzs,tfield,tfield,fftW_FORWARD,fftW_PATIENT)
 write(0,*) "planF computed"
    call sfftw_plan_dft(planI,3,nzs,tfield,tfield,fftw_BACKWARD,fftW_PATIENT)
 write(0,*) "planI computed"

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

stat = 0

  end function fftvar_init
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 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
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 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(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
 cig_out = sog_ifft(Rcig)

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

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

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

 tfield=cmplx(cig,0.)
 call sfftw_execute(planF)
 Ccig=tfield*sqrt(1./(1.*nzs(1)*nzs(2)*nzs(3)));

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

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

 allocate(cig(nz,nh,nx))

 tfield=Ccig;
 call sfftw_execute(planI)
 cig=real(tfield(1:nz,1:nh,1:nx))*sqrt(1./(1.*nzs(1)*nzs(2)*nzs(3)));

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

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

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

k=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=1.E+37
             else if(kx%all(ix).gt.0.) then
                talfa=-1.E+37
             else
                talfa=-kx%d/kz%d
             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
                tteta_s=-1.E+37
             else if(kh%all(ih).gt.0.) then
                tteta_s=1.E+37
             else
		tteta_s=-kh%d/kz%d
             end if
          else
             tteta_s=-kh%all(ih)/kz%all(iz)
          end if

	  if (abs(1.+talfa*tteta_s).lt.(1.E-37)) then
	     if(tteta_s-talfa.lt.0.) then
	       tgama=-1.E+37
	     else if(tteta_s-talfa.gt.0.) then
	       tgama=1.E+37
	     else
	       tgama=0.
	     end if
	  else
             tgama=(tteta_s-talfa)/(1.+talfa*tteta_s)
	  end if

	  khrot=-kz%all(iz)*tgama
	  ikhrot=int(khrot/kh%d)+1
          ilhrot=ikhrot+1
          if (ikhrot < 0) ikhrot = kh%n + ikhrot
          if (ilhrot < 0) ilhrot = kh%n + ilhrot


	  fwrot=khrot/kh%d-ikhrot-1
	  gwrot=1.-fwrot

  	  if (ikhrot.ge.1 .and. ikhrot.le.kh%n .and. ilhrot.ge.1 .and. ilhrot.le.kh%n) then
	     k=k+1
 	     Scig(iz,ih,ix)=gwrot*Ccig(iz,ikhrot,ix)+fwrot*Ccig(iz,ilhrot,ix)+Scig(iz,ih,ix)
! 	     Scig(iz,ih,ix)=Ccig(iz,ikhrot,ix)+Scig(iz,ih,ix)
	  end if

       end do
    end do
 end do
 write(0,*) k

 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,khrot

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

k=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=1.E+37
          else if(kx%all(ix).gt.0.) then
             talfa=-1.E+37
          else
             talfa=-kx%d/kz%d
          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
                tteta_g=-1.E+37
             else if(kh%all(ih).gt.0.) then
                tteta_g=1.E+37
             else
		tteta_g=kh%d/kz%d
             end if
          else
             tteta_g=kh%all(ih)/kz%all(iz)
          end if

	  if (abs(1.+talfa*tteta_g).lt.(1.E-37)) then
	     if((talfa-tteta_g).lt.0.) then
	       tgama=-1.E+37
	     else if((talfa-tteta_g).gt.0.) then
	       tgama=1.E+37
	     else
	       tgama=0.
	     end if
	  else
             tgama=(talfa-tteta_g)/(1.+talfa*tteta_g)
	  end if

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

	  fwrot=(khrot-kh%o)/kh%d-ikhrot-1
	  gwrot=1.-fwrot

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

       end do
    end do
 end do
 write(0,*) k

 end function rotrec_geodip
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
end module sogrot_mod3z
