/* -*- C++ -*- */
#ifndef __fft_header_file__
#define __fft_header_file__

#include <valarray>
#include <complex>

#include "types.h"
#include "Array2D.h"

namespace fft
{


   /**
    * Fast Fourier Transform for 2D data
    * 
    * @param arr Input and output data
    * @param w Width of 2D data (height is computed from input array size)
    * @param forward Type of FFT (true for forward, false for inverse)
    */
   void fft2(Array2D<Complex>& arr);
   void ifft2(Array2D<Complex>& arr);

   /**
    * Shift ouput of 2D FFT, so DC componenet is in center.
    * 
    * @param Input and ouput data
    * @param w Width of data (height is computed from input array size)
    */
   template <typename T>
   void fftshift2(Array2D<T>& arr)
   {
      u32 w = arr.W();
      u32 h = arr.H();
      u32 size = w * h;
      u32 w2 = w/2;
      u32 h2 = h/2;
		
      Array2D<T> ret(w,h);
		
      for (u32 i = 0; i < size; i++)
      {
	 u32 xo = i % w;
	 u32 yo = i / w;
	 u32 xn = (xo + w2) % w;
	 u32 yn = (yo + h2) % h;
	 u32 j = yn * w + xn;
	 ret[j] = arr[i];
      }
		
      arr = ret;
   }

   /**
    * Inverse Shift ouput of 2D FFT, so DC componenet is in top left corner
    * 
    * @param Input and ouput data
    * @param w Width of data (height is computed from input array size)
    */
   template <typename T>
   void ifftshift2(Array2D<T>& arr)
   {
      u32 w = arr.W();
      u32 h = arr.H();
      u32 size = w * h;
      u32 w2 = w/2;
      u32 h2 = h/2;

      Array2D<T> ret(w,h);
		
      for (u32 i = 0; i < size; i++)
      {
	 u32 xo = i % w;
	 u32 yo = i / w;
	 s32 xn = xo - w2;
	 if (xn < 0) xn += w; 
	 s32 yn = yo - h2;
	 if (yn < 0) yn += h;
	 u32 j = static_cast<u32>(yn * w + xn);
	 ret[j] = arr[i];
      }
		
      arr = ret;
   }
;
}

#endif //__fft_header_file__

