/***************************************************************************
 *            rtfft.h
 *
 *  Mon Feb 22 22:14:17 2010
 *  Copyright  2010  Ralph Tandetzky
 *
 *    This header file defines some useful FFT algorithms for C++.
 *  It is not meant to contain the fastest FFT algorithms ever,
 *  but is rather written for educational purposes. However, the
 *  speed is quite acceptable.
 *
 *
 *  The following functions are contained in this file:
 *
 *  DFT(In, Out, N)        calculates Discrete Fourier Transform.
 *  FFT(In, Out, N)        calculates DFT in O(N*log(N)) time.
 *  FFT_Radix2CooleyTukey(In,Out,Work,N)
 *                        Auxiliary function, that calculates DFT in O(N*log(N))
 *                        for N being a power of 2 using a preallocated working
 *                        array.
 *  FFT_Bluestein(In,Out,Work,N)
 *                        Auxiliary function, that calculates DFT in O(N*log(N))
 *                        for arbitrary N. It is multiple times slower than
 *                        FFT_Radix2CooleyTukey(...).
 *
 ****************************************************************************/

/*
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Library General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor Boston, MA 02110-1301,
 * USA
 */

#ifndef __RT_FFT_H
#define __RT_FFT_H

// STL include files
#include<complex>
#include<vector>

using namespace std;

/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *\
    Function:
          DFT

      Purpose:
          Calculate the Discrete Fourier-Transform in O(N2) time, where
          N is the number of array elements given by [N].

      Arguments:
          acIn:
              Input array of size [N].
          aOut:
              Output array. The space must be allocated. The algorithm works
              in-place as well as out-of-place, i.e. [aOut] may be
              equal to [acIn]. However, they may not overlap in any other way.
          N:
              Size of the arrays.

      Return value:
          None. The results will be stored in [aOut].
\* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
template<class T>
void DFT(const complex<T>  *acIn, complex<T>  *aOut, size_t N)
{
    if ( acIn != aOut ) // out-of-place?
    {
        for ( size_t k=0; k<N; ++k )
        {
            complex<T>  Sum(0);
            for (size_t n=0; n<N; ++n )
                Sum+=acIn[n]*polar((T)1,(T)(-2*3.1415926535897932)*n*k/N);
            aOut[k]=Sum;
        }
    }
    else // in-place?
    {
        // Copy the contents of the input array into a temporary array and
        // execute the out-of-place-DFT for that array
        DFT(&vector<complex<T>  >(acIn, acIn+N).front(), aOut, N );
    }
}



/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *\
    Function:
          FFT

      Purpose:
          Calculate the Discrete Fourier-Transform in O(N*log(N)) time, where
          N is the number of array elements given by [N].

      Arguments:
          acIn:
              Input array of size [N].
          aOut:
              Output array. The space must be allocated. The algorithm works
              in-place as well as out-of-place, i.e. [aOut] may be
              equal to [acIn]. However, they may not overlap in any other way.
          N:
              Size of the arrays.

      Return value:
          None. The results will be stored in [aOut].
\* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
template<class T>
void FFT(const complex<T>  *acIn, complex<T>  *aOut, size_t N)
{
    if ( ( N&  (N-1) ) == 0 ) // size is a power of 2?
        FFT_Radix2CooleyTukey(acIn, aOut,&vector<complex<T>  >(N).front(), N);
    else // not a power of 2?
        FFT_Bluestein(acIn, aOut, N);
}



/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *\
    Function:
          FFT_Radix2CooleyTukey

      Purpose:
          Calculate the Discrete Fourier-Transform in O(N*log(N)) time using the
          Cooley-Tukey-Algorithm, where N is the number of array elements given
          by [N]. N has to be a power of 2 or zero.

      Arguments:
          acIn:
              Input array of size [N].
          aOut:
              Output array. The space must be allocated.  The algorithm works
              in-place as well as out-of-place, i.e. [aOut] may be
              equal to [acIn]. However, they may not overlap in any other way.
          aWork:
              Array of size [N], that is used to perform the calculations.
          N:
              Size of the arrays. Must be a power of 2 or zero.

      Return value:
          None. The results will be stored in [aOut].
\* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
template<class T>
void FFT_Radix2CooleyTukey(const complex<T>  *acIn, complex<T>  *aOut,
                           complex<T>  *aWork, size_t N)
{
    if (N == 1)
        aOut[0]=acIn[0];
    else
    {
        size_t N_2=N>>1;

        // Seperate entries with odd and even indices
        for ( size_t i=0; i<N_2; ++i )
        {
            aWork[i    ]=acIn[2*i  ];
            aWork[i+N_2]=acIn[2*i+1];
        }

        // Calculate DFTs for subarrays
        if ( N_2>  1 )
        {
            FFT_Radix2CooleyTukey(aWork    , aWork    , aOut, N_2);
            FFT_Radix2CooleyTukey(aWork+N_2, aWork+N_2, aOut, N_2);
        }

        // Put the two arrays together in the proper way
        complex<T>  MultIncrement( polar((T)1,(T)-3.1415926535798932/N_2) );
        complex<T>  TwiddleFactor( (T)1 );
        for ( size_t i=0; i<N_2; ++i )
        {
            complex<T>  A  = aWork[i];
            complex<T>  B  = TwiddleFactor*aWork[i+N_2];
            aOut[i    ]   = A + B;
            aOut[i+N_2]   = A - B;
            TwiddleFactor*= MultIncrement;
        }
    }
}



/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *\
    Function:
          FFT_Bluestein

      Purpose:
          Calculate the Discrete Fourier-Transform in O(N*log(N)) time using the
          Bluestein-Algorithm, where N is the number of array elements given
          by [N]. N can be any positive integer, however, the algorithm takes
          much more time than the Cooley-Tukey-Algorithm.

      Arguments:
          acIn:
              Input array of size [N].
          aOut:
              Output array. The space must be allocated. The array [aOut]
              may be equal to [acIn].
          N:
              Size of the arrays. Any positive integer.

      Return value:
          None. The results will be stored in [aOut].
\* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
template<class T>
void FFT_Bluestein(const complex<T>  *acIn, complex<T>  *aOut, size_t N)
{
    // BUG WORKAROUND: The actual algorithm doesn't work for odd N. It yields
    // uncorrect results for unknown reasons. Therefore we double the size to
    // keep the O(N*log(N)) time complexity and still get correct results.
    if (N % 2 == 1) // N odd?
    {
        vector<complex<T>  >  vInOut2( 2*N, complex<T>(0) );
        copy( acIn, acIn+N, vInOut2.begin() ); // copy the original input array
        FFT_Bluestein(&vInOut2.front(),&vInOut2.front(), 2*N );
        for ( size_t i=0; i<N; i++ )
            aOut[i] = vInOut2[2*i];
    }
    else // N even?
    {
        // Calculate M as the smallest power of 2 that is at least 2*N, where
        // N is the vector size.
        size_t M=N;
        M--; M|=M>>1; M|=M>>2; M|=M>>4; M|=M>>8; M|=M>>16; M++; M<<=1;

        // Calculate the zero-padded a and b.
        vector<complex<T>  >  a(M,0), b(M,0);
        const T MinAngle=(T)3.1415926535897932/N;
        for ( size_t i=0; i<N; ++i )
        {
            b[i]=polar((T)1,i*i*MinAngle);
            a[i]=conj(b[i])*acIn[i];
        }
        vector<complex<T>  >  bNonZeroPadded(b.begin(), b.begin()+N);

        // Apply Cooley-Tukey-Algorithm to a and b
        vector<complex<T>  >  Work(M);
        FFT_Radix2CooleyTukey(&a.front(),&a.front(),&Work.front(),M);
        FFT_Radix2CooleyTukey(&b.front(),&b.front(),&Work.front(),M);

        // Multiply F(a) and F(b) pointwise, assign to a
        for ( size_t i=0; i<M; ++i )
            a[i] *= b[i];

        // back-transform this product. Calculate convolution of
        // non-zero-padded a and b.
        FFT_Radix2CooleyTukey(&a.front(),&a.front(),&Work.front(),M);
        // Now a contains the convolution the zero-padded a and b times M
        // backwards.
        const T Frac1M=(T)1/M;
        aOut[0] = (a[0]+a[M-N])*Frac1M;
        for ( size_t i=1; i<N; ++i )
            aOut[i] = conj(bNonZeroPadded[i])*(a[M-i]+a[M-i-N])*Frac1M;
    }
}


#endif //__RT_FFT_H