#include "fft.h"
#include <iostream>
#include <cmath>

using namespace std;
namespace fft
{
//typedef complex<f64> Complex;
//typedef valarray<Complex> Array;

   template <typename T>
   class ArrayIterator
   {
      public:
	 ArrayIterator(Array2D<T>& parent, u32 pos)
	    : parent(parent), pos(pos)
	 {}
	 virtual ~ArrayIterator(){}

	 virtual T& operator[](u32 n) = 0;
	 virtual u32 size() const = 0;

	 void operator/=(T rhs)
	 {
	    for (u32 i = 0; i < size(); i++)
	    {
	       (*this)[i] /= rhs;
	    }
	 }

      protected:
	 Array2D<T>& parent;
	 u32 pos;
   };


   template <typename T>
   class RowIterator : public ArrayIterator<T>
   {
      public:
	 RowIterator(Array2D<T>& parent, u32 pos)
	    : ArrayIterator<T>(parent, pos){}

	 virtual T& operator[](u32 i)
	 {
	    return this->parent.at(i,this->pos);
	 }

	 virtual u32 size() const
	 {
	    return this->parent.W();
	 }
   };


   template <typename T>
   class ColumnIterator : public ArrayIterator<T>
   {
      public:
	 ColumnIterator(Array2D<T>& parent, u32 pos)
	    : ArrayIterator<T>(parent, pos){}

	 virtual T& operator[](u32 i)
	 {
	    return this->parent.at(this->pos,i);
	 }

	 virtual u32 size() const
	 {
	    return this->parent.H();
	 }
   };

   template <typename T>
   void fft(ArrayIterator<T>& arr, bool forward)
   {
      size_t n = arr.size();

		
      // Do the bit reversal 
      u32 i2 = n >> 1;
      u32 j = 0;
      for (u32 i=0; i < n-1; i++) 
      {
	 if (i < j) 
	 {
	    Complex tmp = arr[i];
	    arr[i] = arr[j];
	    arr[j] = tmp;
	 }
	 u32 k = i2;
	 while (k <= j) {
	    j -= k;
	    k >>= 1;
	 }
	 j += k;
      }
		
      // do computation
		
		
		
      //Complex c(-1, 0);
      u32 l1;
      u32 l2 = 1;
      u32 m = static_cast<u32>(log(n)/log(2));
      //cout << "M = " <<  m << endl;

      // sign of exponent
      f64 sign = (forward) ? -1.0 : 1.0;
      // base component of exponent of C
      f64 cExpBase = sign * M_PI;

      // initial 
      //Complex c = -1.0;
      // loop for each level
      for (u32 l = 0; l < m; l++)  
      {
	 // count of sub frequency spectras
	 l1 = l2;
	 // sub DFT stride
	 l2 <<= 1;
			
	 // W ^ 0
	 Complex w(1,0);
	 // W ^ 1
	 Complex c(cos(cExpBase/l1), sin(cExpBase/l1)); //= exp(Complex(0, cExpBase / l1));
			
	 // loop for each sub DFT element
	 for (u32 j = 0; j < l1; j++) 
	 {
				
	    // loop for each butterfly (sub DFT)
	    for (u32 i = j; i < n; i += l2) 
	    {
	       u32 i2 = i + l1;
	       Complex t = w * arr[i2];
	       arr[i2] = arr[i] - t;
	       arr[i] += t;
	    }

	    // compute new w (W^j)
	    w *= c;
	 }
	 /*f64 cr = c.real();
	   c.imag() = sign * sqrt((1.0 - cr) / 2.0);
	   c.real() = sqrt((1.0 + cr) / 2.0);*/
			
      }
		
      // normalize
      if (!forward)
      {
	 arr /= static_cast<f64>(n);
      }
   }



   void fft2(Array2D<Complex>& arr)
   {
      for (u32 i = 0; i < arr.W(); i++)
      {
	 ColumnIterator<Complex> it(arr, i);
	 fft(it, true);
      }

      for (u32 j = 0; j < arr.H(); j++)
      {
	 RowIterator<Complex> it(arr, j);
	 fft(it, true);
      }
   }



   void ifft2(Array2D<Complex>& arr)
   {

      for (u32 i = 0; i < arr.W(); i++)
      {
	 ColumnIterator<Complex> it(arr, i);
	 fft(it, false);
      }

      for (u32 j = 0; j < arr.H(); j++)
      {
	 RowIterator<Complex> it(arr, j);
	 fft(it, false);
      }
      //arr.flipV();
   }

	
	
	
}


