module teste02

  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 :: Rot
  
  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
    real,allocatable	   :: k_row1(:)

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

    stat=1
!!!
   allocate (k_row1(nh),kh%all(nh))

        k_row1=0; kh%all=0.
        kh%n=nh
        kh%o=0.
        kh%d=2.*pi/(dh*FLOAT(nh))
        do ikh=1,nh
           if((ikh - 1) <= float(nh/2)) then
              kh%all(ikh) =  kh%o+kh%d*(ikh - 1)
           else
              kh%all(ikh) = -kh%o+kh%d*((ikh - 1)-float(nh/2))
           end if
        end do

! remap k_hdip to facilitate interpolation
        k_row1(:)=kh%all(:)
        do ikh=1,int(kh%n/2)+1
          kh%all(ikh+int(kh%n/2)-1)=k_row1(ikh)
        end do
        do ikh=int(kh%n/2)+2,kh%n
          kh%all(ikh-int(kh%n/2)-1)=k_row1(ikh)
        end do
!!!
    kx%n = nx
    kx%d = 2*pi/(kx%n*dx)
    kx%o = 0.
    allocate(kx%all(kx%n)); kx%all=0.
    do ikx=1,nx
       if((ikx-1) <= float(nx/2)) then
          kx%all(ikx) =  kx%o+kx%d*(ikx-1)
       else
          kx%all(ikx) = -kx%o+kx%d*((ikx-1)-kx%n)
       end if
    end do
!!!
    kz%n = nz
    kz%d = 2*pi/(kz%n*dz)
    kz%o = 0.
    allocate(kz%all(kz%n)); kz%all=0.
    do ikz=1,nz
       if((ikz-1) <= float(nz/2)) then
          kz%all(ikz) =  kz%o+kz%d* (ikz-1)
       else
          kz%all(ikz) = -kz%o+kz%d*((ikz-1)-nz)
       end if
    end do
!!!
    fft_1d_idx = init_fft_1d(kx%n)
    fft_1d_idz = init_fft_1d(kz%n)
    fft_2d_idxh = init_fft_2d(kx%n,kh%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 Rot(cube,sign_tan) 
    
    complex,allocatable					:: k_cube(:,:,:),k_temp(:),k_row(:)
    complex,allocatable					:: temp_z(:),temp_x_hdip(:,:)
    real						:: cube(:,:,:)
    integer status

    logical inside_k_hdip
    integer i_z,i_mx,i_hdip,i_k_hdip,sign_tan
    real geol_dip,off_dip,inp_dip,shift
    real*8 tan_geol_dip,tan_off_dip,tan_inp_dip,tan_out_dip
    real inv_d_khdip,k_hdip_inp,k_hdip_out,w_hdip

    inv_d_khdip=1./kh%d

 allocate(k_cube(kz%n,kx%n,kh%n),k_temp(kh%n),k_row(kh%n))
 allocate(temp_z(kz%n),temp_x_hdip(kx%n,kh%n))   

! zeroes k_cube
    k_cube=cmplx(0.,0.)
    k_temp=cmplx(0.,0.)
    k_row=cmplx(0.,0.)

  ! FFT along z
    do i_mx=1,kx%n
      do i_hdip=1,kh%n
        temp_z(1:kz%n)=cmplx(cube(1:kz%n,i_mx,i_hdip),0.)
        call fft_1d_data(fft_1d_idz,-1,kz%n,temp_z(:))
        do i_z=1,int(kz%n/2)+1
          k_cube(i_z,i_mx,i_hdip)=temp_z(i_z)
        end do
      end do 
    end do

  ! FFT along mx and hdip
    do i_z=1,int(kz%n/2)+1
      temp_x_hdip(:,:)=k_cube(i_z,:,:)
      call fft_2d_data(fft_2d_idxh,-1,kx%n,kh%n,temp_x_hdip(:,:))
      k_cube(i_z,:,:)=temp_x_hdip(:,:)
    end do 

    do i_mx=1,kx%n
      do i_z=1,int(kz%n/2)+1
        if(i_z /= 1) then
          tan_geol_dip=-kx%all(i_mx)/kz%all(i_z)
        else
          tan_geol_dip=sign_tan*sign(huge(tan_geol_dip),-kx%all(i_mx))
        end if

        do i_hdip=1,int(kh%n/2)+1
          k_temp(i_hdip+int(kh%n/2)-1)=k_cube(i_z,i_mx,i_hdip)
        end do
        do i_hdip=int(kh%n/2)+2,kh%n
          k_temp(i_hdip-int(kh%n/2)-1)=k_cube(i_z,i_mx,i_hdip)
        end do

        do i_hdip=1,kh%n
          shift=kh%all(i_hdip)*kh%o
          k_temp(i_hdip)=k_temp(i_hdip)*cexp(cmplx(0,shift))
        end do

          do i_hdip=1,kh%n
            if(i_z /= 1) then
              tan_off_dip=-kh%all(i_hdip)/kz%all(i_z)
! Notice opposite sign in adding tangents
              if (abs(1+sign_tan*tan_off_dip*tan_geol_dip)>tiny(tan_off_dip)) &
                   then
                tan_inp_dip= &
                  (tan_geol_dip - sign_tan*tan_off_dip)/&
                  (1+sign_tan*tan_off_dip*tan_geol_dip)
  
                k_hdip_inp=-kz%all(i_z)*tan_inp_dip
              else
                k_hdip_inp=sign(kh%all(i_hdip),-kx%all(i_mx))
              end if
            else
                k_hdip_inp=sign(kh%all(i_hdip),-kx%all(i_mx))
            end if

            i_k_hdip=floor(k_hdip_inp*inv_d_khdip+float(int(kh%n/2)))
            inside_k_hdip=((i_k_hdip >= 1) .and.(i_k_hdip < kh%n))
  
            w_hdip=(k_hdip_inp*inv_d_khdip)+int(kh%n/2) - i_k_hdip
  
            if(inside_k_hdip) then
              k_row(i_hdip)= &
                           (1.-w_hdip)*k_temp(i_k_hdip  ) + &
                           (   w_hdip)*k_temp(i_k_hdip+1)
            else
              if(i_k_hdip == kh%n) then
                k_row(i_hdip)= &
                             (1.-w_hdip)*k_temp(i_k_hdip  ) + &
                             (   w_hdip)*k_temp(1         )
              else
                if(i_k_hdip == 0) then
                  k_row(i_hdip)= &
                             (1.-w_hdip)*k_temp(kh%n) + &
                             (   w_hdip)*k_temp(1         )
                end if
              end if
            end if
          end do 

        do i_hdip=1,kh%n
          shift=-kh%all(i_hdip)*kh%o
          k_row(i_hdip)=k_row(i_hdip)*cexp(cmplx(0,shift))
        end do

        do i_hdip=1,int(kh%n/2)+1
          k_cube(i_z,i_mx,i_hdip)=k_row(i_hdip+int(kh%n/2)-1)
        end do
        do i_hdip=int(kh%n/2)+2,kh%n
          k_cube(i_z,i_mx,i_hdip)=k_row(i_hdip-int(kh%n/2)-1)
        end do

      end do 
    end do 


! INV FFT along mx and hdip

    do i_z=1,int(kz%n/2)+1
      temp_x_hdip(:,:)=k_cube(i_z,:,:)
      call fft_2d_data(fft_2d_idxh,1,kx%n,kh%n,temp_x_hdip(:,:))
      k_cube(i_z,:,:)=temp_x_hdip(:,:)
    end do 

! INV FFT along z
    do i_mx=1,kx%n
      do i_hdip=1,kh%n
        do i_z=1,int(kz%n/2)+1
          temp_z(i_z)=k_cube(i_z,i_mx,i_hdip)
        end do
        do i_z=int(kz%n/2)+2,kz%n
          temp_z(i_z)= conjg(k_cube(kz%n-i_z+2,i_mx,i_hdip))
        end do

        call fft_1d_data(fft_1d_idz,1,kz%n,temp_z(:))
        cube(1:kz%n,i_mx,i_hdip)=real(temp_z(1:kz%n))

      end do 
    end do 

  return	

  end subroutine Rot
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
end module teste02
