/*
 * fftutil.c
 *
 *  Created on: Jan 5, 2018
 *      Author: deanm
 */

#include "Adafruit_ZeroFFT.h"

/*
 * @brief  In-place bit reversal function.
 * @param[in, out] *pSrc        points to the in-place buffer of Q15 data type.
 * @param[in]      fftLen       length of the FFT.
 * @param[in]      bitRevFactor bit reversal modifier that supports different
 * size FFTs with the same bit reversal table
 * @param[in]      *pBitRevTab  points to bit reversal table.
 * @return none.
 */

static q15_t ALIGN4 scratchData[ZERO_FFT_MAX];

void arm_bitreversal_q15(q15_t *pSrc16, uint32_t fftLen, uint16_t bitRevFactor,
                         uint16_t *pBitRevTab) {
  q31_t *pSrc = (q31_t *)pSrc16;
  q31_t in;
  uint32_t fftLenBy2, fftLenBy2p1;
  uint32_t i, j;

  /*  Initializations */
  j = 0u;
  fftLenBy2 = fftLen / 2u;
  fftLenBy2p1 = (fftLen / 2u) + 1u;

  /* Bit Reversal Implementation */
  for (i = 0u; i <= (fftLenBy2 - 2u); i += 2u) {
    if (i < j) {
      /*  pSrc[i] <-> pSrc[j]; */
      /*  pSrc[i+1u] <-> pSrc[j+1u] */
      in = pSrc[i];
      pSrc[i] = pSrc[j];
      pSrc[j] = in;

      /*  pSrc[i + fftLenBy2p1] <-> pSrc[j + fftLenBy2p1];  */
      /*  pSrc[i + fftLenBy2p1+1u] <-> pSrc[j + fftLenBy2p1+1u] */
      in = pSrc[i + fftLenBy2p1];
      pSrc[i + fftLenBy2p1] = pSrc[j + fftLenBy2p1];
      pSrc[j + fftLenBy2p1] = in;
    }

    /*  pSrc[i+1u] <-> pSrc[j+fftLenBy2];         */
    /*  pSrc[i+2] <-> pSrc[j+fftLenBy2+1u]  */
    in = pSrc[i + 1u];
    pSrc[i + 1u] = pSrc[j + fftLenBy2];
    pSrc[j + fftLenBy2] = in;

    /*  Reading the index for the bit reversal */
    j = *pBitRevTab;

    /*  Updating the bit reversal index depending on the fft length  */
    pBitRevTab += bitRevFactor;
  }
}

void arm_radix2_butterfly_q15(q15_t *pSrc, uint32_t fftLen, q15_t *pCoef,
                              uint16_t twidCoefModifier) {
  int i, j, k, l;
  int n1, n2, ia;
  q15_t xt, yt, cosVal, sinVal;

  // N = fftLen;
  n2 = fftLen;

  n1 = n2;
  n2 = n2 >> 1;
  ia = 0;

  // loop for groups
  for (j = 0; j < n2; j++) {
    cosVal = pCoef[ia * 2];
    sinVal = pCoef[(ia * 2) + 1];
    ia = ia + twidCoefModifier;

    // loop for butterfly
    for (i = j; i < fftLen; i += n1) {
      l = i + n2;
      xt = (pSrc[2 * i] >> 2u) - (pSrc[2 * l] >> 2u);
      pSrc[2 * i] = ((pSrc[2 * i] >> 2u) + (pSrc[2 * l] >> 2u)) >> 1u;

      yt = (pSrc[2 * i + 1] >> 2u) - (pSrc[2 * l + 1] >> 2u);
      pSrc[2 * i + 1] =
          ((pSrc[2 * l + 1] >> 2u) + (pSrc[2 * i + 1] >> 2u)) >> 1u;

      pSrc[2u * l] = (((int16_t)(((q31_t)xt * cosVal) >> 16)) +
                      ((int16_t)(((q31_t)yt * sinVal) >> 16)));

      pSrc[2u * l + 1u] = (((int16_t)(((q31_t)yt * cosVal) >> 16)) -
                           ((int16_t)(((q31_t)xt * sinVal) >> 16)));

    } // butterfly loop end

  } // groups loop end

  twidCoefModifier = twidCoefModifier << 1u;

  // loop for stage
  for (k = fftLen / 2; k > 2; k = k >> 1) {
    n1 = n2;
    n2 = n2 >> 1;
    ia = 0;

    // loop for groups
    for (j = 0; j < n2; j++) {
      cosVal = pCoef[ia * 2];
      sinVal = pCoef[(ia * 2) + 1];
      ia = ia + twidCoefModifier;

      // loop for butterfly
      for (i = j; i < fftLen; i += n1) {
        l = i + n2;
        xt = pSrc[2 * i] - pSrc[2 * l];
        pSrc[2 * i] = (pSrc[2 * i] + pSrc[2 * l]) >> 1u;

        yt = pSrc[2 * i + 1] - pSrc[2 * l + 1];
        pSrc[2 * i + 1] = (pSrc[2 * l + 1] + pSrc[2 * i + 1]) >> 1u;

        pSrc[2u * l] = (((int16_t)(((q31_t)xt * cosVal) >> 16)) +
                        ((int16_t)(((q31_t)yt * sinVal) >> 16)));

        pSrc[2u * l + 1u] = (((int16_t)(((q31_t)yt * cosVal) >> 16)) -
                             ((int16_t)(((q31_t)xt * sinVal) >> 16)));

      } // butterfly loop end

    } // groups loop end

    twidCoefModifier = twidCoefModifier << 1u;
  } // stages loop end

  n1 = n2;
  n2 = n2 >> 1;
  ia = 0;

  // loop for groups
  for (j = 0; j < n2; j++) {
    cosVal = pCoef[ia * 2];
    sinVal = pCoef[(ia * 2) + 1];

    ia = ia + twidCoefModifier;

    // loop for butterfly
    for (i = j; i < fftLen; i += n1) {
      l = i + n2;
      xt = pSrc[2 * i] - pSrc[2 * l];
      pSrc[2 * i] = (pSrc[2 * i] + pSrc[2 * l]);

      yt = pSrc[2 * i + 1] - pSrc[2 * l + 1];
      pSrc[2 * i + 1] = (pSrc[2 * l + 1] + pSrc[2 * i + 1]);

      pSrc[2u * l] = xt;

      pSrc[2u * l + 1u] = yt;

    } // butterfly loop end

  } // groups loop end

  twidCoefModifier = twidCoefModifier << 1u;
}

static inline void applyWindow(q15_t *src, const q15_t *window, uint16_t len) {
  while (len--) {
    int32_t val = *src * *window++;
    *src++ = val >> 15;
  }
}

int ZeroFFT(q15_t *source, uint16_t length) {
  uint16_t twidCoefModifier;
  uint16_t bitRevFactor;
  uint16_t *pBitRevTable;

  q15_t *pSrc = source;

  switch (length) {

#if ZERO_FFT_MAX == 8192
  case 4096u:
    /*  Initializations of structure parameters for 4096 point FFT */

    /*  Initialise the twiddle coef modifier value */
    twidCoefModifier = 1u;
    /*  Initialise the bit reversal table modifier */
    bitRevFactor = 1u;
    /*  Initialise the bit reversal table pointer */
    pBitRevTable = (uint16_t *)armBitRevTable;

    applyWindow(source, window_hanning_4096, 4096);

    break;
#endif

#if ZERO_FFT_MAX >= 4096

  case 2048u:
    /*  Initializations of structure parameters for 2048 point FFT */

    /*  Initialise the twiddle coef modifier value */
    twidCoefModifier = 2u;
    /*  Initialise the bit reversal table modifier */
    bitRevFactor = 2u;
    /*  Initialise the bit reversal table pointer */
    pBitRevTable = (uint16_t *)&armBitRevTable[1];

    applyWindow(source, window_hanning_2048, 2048);

    break;

#endif

#if ZERO_FFT_MAX >= 2048
  case 1024u:
    /*  Initializations of structure parameters for 1024 point FFT */
    twidCoefModifier = 4u;
    bitRevFactor = 4u;
    pBitRevTable = (uint16_t *)&armBitRevTable[3];

    applyWindow(source, window_hanning_1024, 1024);
    break;
#endif

#if ZERO_FFT_MAX >= 1024
  case 512u:
    /*  Initializations of structure parameters for 512 point FFT */
    twidCoefModifier = 8u;
    bitRevFactor = 8u;
    pBitRevTable = (uint16_t *)&armBitRevTable[7];

    applyWindow(source, window_hanning_512, 512);

    break;
#endif

#if ZERO_FFT_MAX >= 512
  case 256u:
    /*  Initializations of structure parameters for 256 point FFT */
    twidCoefModifier = 16u;
    bitRevFactor = 16u;
    pBitRevTable = (uint16_t *)&armBitRevTable[15];

    applyWindow(source, window_hanning_256, 256);

    break;
#endif

  case 128u:
    /*  Initializations of structure parameters for 128 point FFT */
    twidCoefModifier = 32u;
    bitRevFactor = 32u;
    pBitRevTable = (uint16_t *)&armBitRevTable[31];

    applyWindow(source, window_hanning_128, 128);

    break;

  case 64u:
    /*  Initializations of structure parameters for 64 point FFT */
    twidCoefModifier = 64u;
    bitRevFactor = 64u;
    pBitRevTable = (uint16_t *)&armBitRevTable[63];

    applyWindow(source, window_hanning_64, 64);

    break;

  case 32u:
    /*  Initializations of structure parameters for 32 point FFT */
    twidCoefModifier = 128u;
    bitRevFactor = 128u;
    pBitRevTable = (uint16_t *)&armBitRevTable[127];

    applyWindow(source, window_hanning_32, 32);

    break;

  case 16u:
    /*  Initializations of structure parameters for 16 point FFT */
    twidCoefModifier = 256u;
    bitRevFactor = 256u;
    pBitRevTable = (uint16_t *)&armBitRevTable[255];

    applyWindow(source, window_hanning_16, 16);

    break;

  default:
    /*  Reporting argument error if fftSize is not valid value */
    return -1;
    break;
  }

  // split the data
  q15_t *pOut = scratchData;
  for (int i = 0; i < length; i++) {
    *pOut++ = *pSrc++; // real
    *pOut++ = 0;       // imaginary
  }

  arm_radix2_butterfly_q15(scratchData, length, (q15_t *)twiddleCoefQ15,
                           twidCoefModifier);
  arm_bitreversal_q15(scratchData, length, bitRevFactor, pBitRevTable);

  pSrc = source;
  pOut = scratchData;
  for (int i = 0; i < length; i++) {
    q15_t val = *pOut++;
    uint32_t v = abs(val);
    *pSrc++ = v;
    pOut++; // discard imaginary phase val
  }

  return 0;
}