module sogrot_mod3a

  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, sog_fft3, sog_ifft3, rotsou_geodip, rotrec_geodip, check_sum
  
  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+npadh
    kh%d = 2*pi/(kh%n*dh)
    kh%o = -(kh%n/2)*kh%d
!    kh%o = 0.
    allocate(kh%all(kh%n)); kh%all=0.
    do ikh = 1, kh%n
      kh%all(ikh) = kh%o+(ikh-1)*kh%d
    end do
!!!
    kx%n = nx+npadx
    kx%d = 2*pi/(kx%n*dx)
    kx%o = -(kx%n/2)*kx%d
!    kx%o = 0.
    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)*kz%d
!    kz%o = 0.
    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)
    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
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 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))
! allocate(Rcig(kz%n,kh%n,kx%n),cig_out(kz%n,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_ifft3(Rcig)
! cig_out = sog_ifft(Rcig)
! cig_out = cabs(Rcig)

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

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

 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

 allocate(temp(kz%n,kh%n,kx%n))
 temp=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/sqrt(1.*kz%n*kh%n*kx%n)

b=check_sum(cabs(Ccig))
write(0,*) "fftchk1 ", b

 deallocate(temp)

 end function sog_fft

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

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

 allocate(cig(nz,nh,nx))
! allocate(cig(kz%n,kh%n,kx%n))
 allocate(temp(kz%n,kh%n,kx%n))
 temp=0.
 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

b=check_sum(cabs(Ccig))
write(0,*) "ifftchk1 ", b

 deallocate(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


b=check_sum(cabs(Ccig))
write(0,*) "ifftchk2 ", b

 cig=real(Ccig(1:nz,1:nh,1:nx))/sqrt(1.*kz%n*kh%n*kx%n)

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

 complex, allocatable			:: Ccig(:,:,:),temp(:,:,:)
 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)

 call fft_3d_data(fft_3d_idzhx,  1,  kz%n,kh%n,kx%n,Ccig)

 allocate(temp(kz%n,kh%n,kx%n))
 temp=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/sqrt(1.*kz%n*kh%n*kx%n)

 deallocate(temp)

 end function sog_fft3
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 function sog_ifft3(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

 deallocate(temp)

 call fft_3d_data(fft_3d_idzhx,  -1,  kz%n,kh%n,kx%n,Ccig)

 cig=real(Ccig(1:nz,1:nh,1:nx))/sqrt(1.*kz%n*kh%n*kx%n)

 end function sog_ifft3
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 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,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
             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%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
  	     Scig(iz,ih,ix)=gwrot*Ccig(iz,ikhrot,ix)+fwrot*Ccig(iz,ikhrot+1,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_mod3a
