/**********************************************************************

  FFT.cpp

  Dominic Mazzoni

  September 2000

*******************************************************************//*!

\file FFT.cpp
\brief Fast Fourier Transform routines.

  This file contains a few FFT routines, including a real-FFT
  routine that is almost twice as fast as a normal complex FFT,
  and a power spectrum routine when you know you don't care
  about phase information.

  Some of this code was based on a free implementation of an FFT
  by Don Cross, available on the web at:

    http://www.intersrv.com/~dcross/fft.html

  The basic algorithm for his code was based on Numerican Recipes
  in Fortran.  I optimized his code further by reducing array
  accesses, caching the bit reversal table, and eliminating
  float-to-double conversions, and I added the routines to
  calculate a real FFT and a real power spectrum.

*//*******************************************************************/
/*
  Salvo Ventura - November 2006
  Added more window functions:
    * 4: Blackman
    * 5: Blackman-Harris
    * 6: Welch
    * 7: Gaussian(a=2.5)
    * 8: Gaussian(a=3.5)
    * 9: Gaussian(a=4.5)
*/

#include "FFT.h"

#include "Internat.h"
#include "MemoryX.h"

#include <wx/wxcrtvararg.h>
#include <stdlib.h>
#include <math.h>

#include "RealFFTf.h"

using Floats = ArrayOf<float>;
static ArraysOf<int> gFFTBitTable;
static const size_t MaxFastBits = 16;

/* Declare Static functions */
static void InitFFT();

static bool IsPowerOfTwo(size_t x)
{
   if (x < 2)
      return false;

   if (x & (x - 1))             /* Thanks to 'byang' for this cute trick! */
      return false;

   return true;
}

static size_t NumberOfBitsNeeded(size_t PowerOfTwo)
{
   if (PowerOfTwo < 2) {
      wxFprintf(stderr, "Error: FFT called with size %ld\n", PowerOfTwo);
      exit(1);
   }

   size_t i = 0;
   while (PowerOfTwo > 1)
      PowerOfTwo >>= 1, ++i;

   return i;
}

int ReverseBits(size_t index, size_t NumBits)
{
   size_t i, rev;

   for (i = rev = 0; i < NumBits; i++) {
      rev = (rev << 1) | (index & 1);
      index >>= 1;
   }

   return rev;
}

void InitFFT()
{
   gFFTBitTable.reinit(MaxFastBits);

   size_t len = 2;
   for (size_t b = 1; b <= MaxFastBits; b++) {
      auto &array = gFFTBitTable[b - 1];
      array.reinit(len);
      for (size_t i = 0; i < len; i++)
         array[i] = ReverseBits(i, b);

      len <<= 1;
   }
}

void DeinitFFT()
{
   gFFTBitTable.reset();
}

static inline size_t FastReverseBits(size_t i, size_t NumBits)
{
   if (NumBits <= MaxFastBits)
      return gFFTBitTable[NumBits - 1][i];
   else
      return ReverseBits(i, NumBits);
}

/*
 * Complex Fast Fourier Transform
 */

void FFT(size_t NumSamples,
         bool InverseTransform,
         const float *RealIn, const float *ImagIn,
	 float *RealOut, float *ImagOut)
{
   double angle_numerator = 2.0 * M_PI;
   double tr, ti;                /* temp real, temp imaginary */

   if (!IsPowerOfTwo(NumSamples)) {
      wxFprintf(stderr, "%ld is not a power of two\n", NumSamples);
      exit(1);
   }

   if (!gFFTBitTable)
      InitFFT();

   if (!InverseTransform)
      angle_numerator = -angle_numerator;

   /* Number of bits needed to store indices */
   auto NumBits = NumberOfBitsNeeded(NumSamples);

   /*
    **   Do simultaneous data copy and bit-reversal ordering into outputs...
    */

   for (size_t i = 0; i < NumSamples; i++) {
      auto j = FastReverseBits(i, NumBits);
      RealOut[j] = RealIn[i];
      ImagOut[j] = (ImagIn == NULL) ? 0.0 : ImagIn[i];
   }

   /*
    **   Do the FFT itself...
    */

   size_t BlockEnd = 1;
   for (size_t BlockSize = 2; BlockSize <= NumSamples; BlockSize <<= 1) {

      double delta_angle = angle_numerator / (double) BlockSize;

      double sm2 = sin(-2 * delta_angle);
      double sm1 = sin(-delta_angle);
      double cm2 = cos(-2 * delta_angle);
      double cm1 = cos(-delta_angle);
      double w = 2 * cm1;
      double ar0, ar1, ar2, ai0, ai1, ai2;

      for (size_t i = 0; i < NumSamples; i += BlockSize) {
         ar2 = cm2;
         ar1 = cm1;

         ai2 = sm2;
         ai1 = sm1;

         for (size_t j = i, n = 0; n < BlockEnd; j++, n++) {
            ar0 = w * ar1 - ar2;
            ar2 = ar1;
            ar1 = ar0;

            ai0 = w * ai1 - ai2;
            ai2 = ai1;
            ai1 = ai0;

            size_t k = j + BlockEnd;
            tr = ar0 * RealOut[k] - ai0 * ImagOut[k];
            ti = ar0 * ImagOut[k] + ai0 * RealOut[k];

            RealOut[k] = RealOut[j] - tr;
            ImagOut[k] = ImagOut[j] - ti;

            RealOut[j] += tr;
            ImagOut[j] += ti;
         }
      }

      BlockEnd = BlockSize;
   }

   /*
      **   Need to normalize if inverse transform...
    */

   if (InverseTransform) {
      float denom = (float) NumSamples;

      for (size_t i = 0; i < NumSamples; i++) {
         RealOut[i] /= denom;
         ImagOut[i] /= denom;
      }
   }
}

/*
 * Real Fast Fourier Transform
 *
 * This is merely a wrapper of RealFFTf() from RealFFTf.h.
 */

void RealFFT(size_t NumSamples, const float *RealIn, float *RealOut, float *ImagOut)
{
   auto hFFT = GetFFT(NumSamples);
   Floats pFFT{ NumSamples };
   // Copy the data into the processing buffer
   for(size_t i = 0; i < NumSamples; i++)
      pFFT[i] = RealIn[i];

   // Perform the FFT
   RealFFTf(pFFT.get(), hFFT.get());

   // Copy the data into the real and imaginary outputs
   for (size_t i = 1; i<(NumSamples / 2); i++) {
      RealOut[i]=pFFT[hFFT->BitReversed[i]  ];
      ImagOut[i]=pFFT[hFFT->BitReversed[i]+1];
   }
   // Handle the (real-only) DC and Fs/2 bins
   RealOut[0] = pFFT[0];
   RealOut[NumSamples / 2] = pFFT[1];
   ImagOut[0] = ImagOut[NumSamples / 2] = 0;
   // Fill in the upper half using symmetry properties
   for(size_t i = NumSamples / 2 + 1; i < NumSamples; i++) {
      RealOut[i] =  RealOut[NumSamples-i];
      ImagOut[i] = -ImagOut[NumSamples-i];
   }
}

/*
 * InverseRealFFT
 *
 * This function computes the inverse of RealFFT, above.
 * The RealIn and ImagIn is assumed to be conjugate-symmetric
 * and as a result the output is purely real.
 * Only the first half of RealIn and ImagIn are used due to this
 * symmetry assumption.
 *
 * This is merely a wrapper of InverseRealFFTf() from RealFFTf.h.
 */
void InverseRealFFT(size_t NumSamples, const float *RealIn, const float *ImagIn,
		    float *RealOut)
{
   auto hFFT = GetFFT(NumSamples);
   Floats pFFT{ NumSamples };
   // Copy the data into the processing buffer
   for (size_t i = 0; i < (NumSamples / 2); i++)
      pFFT[2*i  ] = RealIn[i];
   if(ImagIn == NULL) {
      for (size_t i = 0; i < (NumSamples / 2); i++)
         pFFT[2*i+1] = 0;
   } else {
      for (size_t i = 0; i < (NumSamples / 2); i++)
         pFFT[2*i+1] = ImagIn[i];
   }
   // Put the fs/2 component in the imaginary part of the DC bin
   pFFT[1] = RealIn[NumSamples / 2];

   // Perform the FFT
   InverseRealFFTf(pFFT.get(), hFFT.get());

   // Copy the data to the (purely real) output buffer
   ReorderToTime(hFFT.get(), pFFT.get(), RealOut);
}

/*
 * PowerSpectrum
 *
 * This function uses RealFFTf() from RealFFTf.h to perform the real
 * FFT computation, and then squares the real and imaginary part of
 * each coefficient, extracting the power and throwing away the phase.
 *
 * For speed, it does not call RealFFT, but duplicates some
 * of its code.
 */

void PowerSpectrum(size_t NumSamples, const float *In, float *Out)
{
   auto hFFT = GetFFT(NumSamples);
   Floats pFFT{ NumSamples };
   // Copy the data into the processing buffer
   for (size_t i = 0; i<NumSamples; i++)
      pFFT[i] = In[i];

   // Perform the FFT
   RealFFTf(pFFT.get(), hFFT.get());

   // Copy the data into the real and imaginary outputs
   for (size_t i = 1; i<NumSamples / 2; i++) {
      Out[i]= (pFFT[hFFT->BitReversed[i]  ]*pFFT[hFFT->BitReversed[i]  ])
         + (pFFT[hFFT->BitReversed[i]+1]*pFFT[hFFT->BitReversed[i]+1]);
   }
   // Handle the (real-only) DC and Fs/2 bins
   Out[0] = pFFT[0]*pFFT[0];
   Out[NumSamples / 2] = pFFT[1]*pFFT[1];
}

/*
 * Windowing Functions
 */

int NumWindowFuncs()
{
   return eWinFuncCount;
}

const TranslatableString WindowFuncName(int whichFunction)
{
   switch (whichFunction) {
   default:
   case eWinFuncRectangular:
      return XO("Rectangular");
   case eWinFuncBartlett:
      /* i18n-hint a proper name */
      return XO("Bartlett");
   case eWinFuncHamming:
      /* i18n-hint a proper name */
      return XO("Hamming");
   case eWinFuncHann:
      /* i18n-hint a proper name */
      return XO("Hann");
   case eWinFuncBlackman:
      /* i18n-hint a proper name */
      return XO("Blackman");
   case eWinFuncBlackmanHarris:
      /* i18n-hint two proper names */
      return XO("Blackman-Harris");
   case eWinFuncWelch:
      /* i18n-hint a proper name */
      return XO("Welch");
   case eWinFuncGaussian25:
      /* i18n-hint a mathematical function named for C. F. Gauss */
      return XO("Gaussian(a=2.5)");
   case eWinFuncGaussian35:
      /* i18n-hint a mathematical function named for C. F. Gauss */
      return XO("Gaussian(a=3.5)");
   case eWinFuncGaussian45:
      /* i18n-hint a mathematical function named for C. F. Gauss */
      return XO("Gaussian(a=4.5)");
   }
}

void NewWindowFunc(int whichFunction, size_t NumSamplesIn, bool extraSample, float *in)
{
   int NumSamples = (int)NumSamplesIn;
   if (extraSample) {
      wxASSERT(NumSamples > 0);
      --NumSamples;
   }
   wxASSERT(NumSamples > 0);

   switch (whichFunction) {
   default:
      wxFprintf(stderr, "FFT::WindowFunc - Invalid window function: %d\n", whichFunction);
      break;
   case eWinFuncRectangular:
      // Multiply all by 1.0f -- do nothing
      break;

   case eWinFuncBartlett:
   {
      // Bartlett (triangular) window
      const int nPairs = (NumSamples - 1) / 2; // whether even or odd NumSamples, this is correct
      const float denom = NumSamples / 2.0f;
      in[0] = 0.0f;
      for (int ii = 1;
           ii <= nPairs; // Yes, <=
           ++ii) {
         const float value = ii / denom;
         in[ii] *= value;
         in[NumSamples - ii] *= value;
      }
      // When NumSamples is even, in[half] should be multiplied by 1.0, so unchanged
      // When odd, the value of 1.0 is not reached
   }
      break;
   case eWinFuncHamming:
   {
      // Hamming
      const double multiplier = 2 * M_PI / NumSamples;
      static const double coeff0 = 0.54, coeff1 = -0.46;
      for (int ii = 0; ii < NumSamples; ++ii)
         in[ii] *= coeff0 + coeff1 * cos(ii * multiplier);
   }
      break;
   case eWinFuncHann:
   {
      // Hann
      const double multiplier = 2 * M_PI / NumSamples;
      static const double coeff0 = 0.5, coeff1 = -0.5;
      for (int ii = 0; ii < NumSamples; ++ii)
         in[ii] *= coeff0 + coeff1 * cos(ii * multiplier);
   }
      break;
   case eWinFuncBlackman:
   {
      // Blackman
      const double multiplier = 2 * M_PI / NumSamples;
      const double multiplier2 = 2 * multiplier;
      static const double coeff0 = 0.42, coeff1 = -0.5, coeff2 = 0.08;
      for (int ii = 0; ii < NumSamples; ++ii)
         in[ii] *= coeff0 + coeff1 * cos(ii * multiplier) + coeff2 * cos(ii * multiplier2);
   }
      break;
   case eWinFuncBlackmanHarris:
   {
      // Blackman-Harris
      const double multiplier = 2 * M_PI / NumSamples;
      const double multiplier2 = 2 * multiplier;
      const double multiplier3 = 3 * multiplier;
      static const double coeff0 = 0.35875, coeff1 = -0.48829, coeff2 = 0.14128, coeff3 = -0.01168;
      for (int ii = 0; ii < NumSamples; ++ii)
         in[ii] *= coeff0 + coeff1 * cos(ii * multiplier) + coeff2 * cos(ii * multiplier2) + coeff3 * cos(ii * multiplier3);
   }
      break;
   case eWinFuncWelch:
   {
      // Welch
      const float N = NumSamples;
      for (int ii = 0; ii < NumSamples; ++ii) {
         const float iOverN = ii / N;
         in[ii] *= 4 * iOverN * (1 - iOverN);
      }
   }
      break;
   case eWinFuncGaussian25:
   {
      // Gaussian (a=2.5)
      // Precalculate some values, and simplify the fmla to try and reduce overhead
      static const double A = -2 * 2.5*2.5;
      const float N = NumSamples;
      for (int ii = 0; ii < NumSamples; ++ii) {
         const float iOverN = ii / N;
         // full
         // in[ii] *= exp(-0.5*(A*((ii-NumSamples/2)/NumSamples/2))*(A*((ii-NumSamples/2)/NumSamples/2)));
         // reduced
         in[ii] *= exp(A * (0.25 + (iOverN * iOverN) - iOverN));
      }
   }
      break;
   case eWinFuncGaussian35:
   {
      // Gaussian (a=3.5)
      static const double A = -2 * 3.5*3.5;
      const float N = NumSamples;
      for (int ii = 0; ii < NumSamples; ++ii) {
         const float iOverN = ii / N;
         in[ii] *= exp(A * (0.25 + (iOverN * iOverN) - iOverN));
      }
   }
      break;
   case eWinFuncGaussian45:
   {
      // Gaussian (a=4.5)
      static const double A = -2 * 4.5*4.5;
      const float N = NumSamples;
      for (int ii = 0; ii < NumSamples; ++ii) {
         const float iOverN = ii / N;
         in[ii] *= exp(A * (0.25 + (iOverN * iOverN) - iOverN));
      }
   }
      break;
   }

   if (extraSample && whichFunction != eWinFuncRectangular) {
      double value = 0.0;
      switch (whichFunction) {
      case eWinFuncHamming:
         value = 0.08;
         break;
      case eWinFuncGaussian25:
         value = exp(-2 * 2.5 * 2.5 * 0.25);
         break;
      case eWinFuncGaussian35:
         value = exp(-2 * 3.5 * 3.5 * 0.25);
         break;
      case eWinFuncGaussian45:
         value = exp(-2 * 4.5 * 4.5 * 0.25);
         break;
      default:
         break;
      }
      in[NumSamples] *= value;
   }
}

// See cautions in FFT.h !
void WindowFunc(int whichFunction, size_t NumSamples, float *in)
{
   bool extraSample = false;
   switch (whichFunction)
   {
   case eWinFuncHamming:
   case eWinFuncHann:
   case eWinFuncBlackman:
   case eWinFuncBlackmanHarris:
      extraSample = true;
      break;
   default:
      break;
   case eWinFuncBartlett:
      // PRL:  Do nothing here either
      // But I want to comment that the old function did this case
      // wrong in the second half of the array, in case NumSamples was odd
      // but I think that never happened, so I am not bothering to preserve that
      break;
   }
   NewWindowFunc(whichFunction, NumSamples, extraSample, in);
}

void DerivativeOfWindowFunc(int whichFunction, size_t NumSamples, bool extraSample, float *in)
{
   if (eWinFuncRectangular == whichFunction)
   {
      // Rectangular
      // There are deltas at the ends
      wxASSERT(NumSamples > 0);
      --NumSamples;
      // in[0] *= 1.0f;
      for (int ii = 1; ii < (int)NumSamples; ++ii)
         in[ii] = 0.0f;
      in[NumSamples] *= -1.0f;
      return;
   }

   if (extraSample) {
      wxASSERT(NumSamples > 0);
      --NumSamples;
   }

   wxASSERT(NumSamples > 0);

   double A;
   switch (whichFunction) {
   case eWinFuncBartlett:
   {
      // Bartlett (triangular) window
      // There are discontinuities in the derivative at the ends, and maybe at the midpoint
      const int nPairs = (NumSamples - 1) / 2; // whether even or odd NumSamples, this is correct
      const float value = 2.0f / NumSamples;
      in[0] *=
         // Average the two limiting values of discontinuous derivative
         value / 2.0f;
      for (int ii = 1;
         ii <= nPairs; // Yes, <=
         ++ii) {
         in[ii] *= value;
         in[NumSamples - ii] *= -value;
      }
      if (NumSamples % 2 == 0)
         // Average the two limiting values of discontinuous derivative
         in[NumSamples / 2] = 0.0f;
      if (extraSample)
         in[NumSamples] *=
         // Average the two limiting values of discontinuous derivative
            -value / 2.0f;
      else
         // Halve the multiplier previously applied
         // Average the two limiting values of discontinuous derivative
         in[NumSamples - 1] *= 0.5f;
   }
      break;
   case eWinFuncHamming:
   {
      // Hamming
      // There are deltas at the ends
      const double multiplier = 2 * M_PI / NumSamples;
      static const double coeff0 = 0.54, coeff1 = -0.46 * multiplier;
      // TODO This code should be more explicit about the precision it intends.
      // For now we get C4305 warnings, truncation from 'const double' to 'float' 
      in[0] *= coeff0;
      if (!extraSample)
         --NumSamples;
      for (int ii = 0; ii < (int)NumSamples; ++ii)
         in[ii] *= - coeff1 * sin(ii * multiplier);
      if (extraSample)
         in[NumSamples] *= - coeff0;
      else
         // slightly different
         in[NumSamples] *= - coeff0 - coeff1 * sin(NumSamples * multiplier);
   }
      break;
   case eWinFuncHann:
   {
      // Hann
      const double multiplier = 2 * M_PI / NumSamples;
      const double coeff1 = -0.5 * multiplier;
      for (int ii = 0; ii < (int)NumSamples; ++ii)
         in[ii] *= - coeff1 * sin(ii * multiplier);
      if (extraSample)
         in[NumSamples] = 0.0f;
   }
      break;
   case eWinFuncBlackman:
   {
      // Blackman
      const double multiplier = 2 * M_PI / NumSamples;
      const double multiplier2 = 2 * multiplier;
      const double coeff1 = -0.5 * multiplier, coeff2 = 0.08 * multiplier2;
      for (int ii = 0; ii < (int)NumSamples; ++ii)
         in[ii] *= - coeff1 * sin(ii * multiplier) - coeff2 * sin(ii * multiplier2);
      if (extraSample)
         in[NumSamples] = 0.0f;
   }
      break;
   case eWinFuncBlackmanHarris:
   {
      // Blackman-Harris
      const double multiplier = 2 * M_PI / NumSamples;
      const double multiplier2 = 2 * multiplier;
      const double multiplier3 = 3 * multiplier;
      const double coeff1 = -0.48829 * multiplier,
         coeff2 = 0.14128 * multiplier2, coeff3 = -0.01168 * multiplier3;
      for (int ii = 0; ii < (int)NumSamples; ++ii)
         in[ii] *= - coeff1 * sin(ii * multiplier) - coeff2 * sin(ii * multiplier2) - coeff3 * sin(ii * multiplier3);
      if (extraSample)
         in[NumSamples] = 0.0f;
   }
      break;
   case eWinFuncWelch:
   {
      // Welch
      const float N = NumSamples;
      const float NN = NumSamples * NumSamples;
      for (int ii = 0; ii < (int)NumSamples; ++ii) {
         in[ii] *= 4 * (N - ii - ii) / NN;
      }
      if (extraSample)
         in[NumSamples] = 0.0f;
      // Average the two limiting values of discontinuous derivative
      in[0] /= 2.0f;
      in[NumSamples - 1] /= 2.0f;
   }
      break;
   case eWinFuncGaussian25:
      // Gaussian (a=2.5)
      A = -2 * 2.5*2.5;
      goto Gaussian;
   case eWinFuncGaussian35:
      // Gaussian (a=3.5)
      A = -2 * 3.5*3.5;
      goto Gaussian;
   case eWinFuncGaussian45:
      // Gaussian (a=4.5)
      A = -2 * 4.5*4.5;
      goto Gaussian;
   Gaussian:
   {
      // Gaussian (a=2.5)
      // There are deltas at the ends
      const float invN = 1.0f / NumSamples;
      const float invNN = invN * invN;
      // Simplify formula from the loop for ii == 0, add term for the delta
      in[0] *= exp(A * 0.25) * (1 - invN);
      if (!extraSample)
         --NumSamples;
      for (int ii = 1; ii < (int)NumSamples; ++ii) {
         const float iOverN = ii * invN;
         in[ii] *= exp(A * (0.25 + (iOverN * iOverN) - iOverN)) * (2 * ii * invNN - invN);
      }
      if (extraSample)
         in[NumSamples] *= exp(A * 0.25) * (invN - 1);
      else {
         // Slightly different
         const float iOverN = NumSamples * invN;
         in[NumSamples] *= exp(A * (0.25 + (iOverN * iOverN) - iOverN)) * (2 * NumSamples * invNN - invN - 1);
      }
   }
      break;
   default:
      wxFprintf(stderr, "FFT::DerivativeOfWindowFunc - Invalid window function: %d\n", whichFunction);
   }
}
