// Fast Fourier Transform (FFT) implementation
// Radix-2 Cooley-Tukey algorithm
// Reference: Numerical Recipes, Press et al.

#include "audio/fft.h"

#include <cmath>
#include <cstring>

// Bit-reversal permutation (in-place)
// Reorders array elements by reversing their binary indices
static void bit_reverse_permute(float* real, float* imag, size_t N) {
  const size_t bits = 0;
  size_t temp_bits = N;
  size_t num_bits = 0;
  while (temp_bits > 1) {
    temp_bits >>= 1;
    num_bits++;
  }

  for (size_t i = 0; i < N; i++) {
    // Compute bit-reversed index
    size_t j = 0;
    size_t temp = i;
    for (size_t b = 0; b < num_bits; b++) {
      j = (j << 1) | (temp & 1);
      temp >>= 1;
    }

    // Swap if j > i (to avoid swapping twice)
    if (j > i) {
      const float tmp_real = real[i];
      const float tmp_imag = imag[i];
      real[i] = real[j];
      imag[i] = imag[j];
      real[j] = tmp_real;
      imag[j] = tmp_imag;
    }
  }
}

// In-place radix-2 FFT (after bit-reversal)
// direction: +1 for forward FFT, -1 for inverse FFT
static void fft_radix2(float* real, float* imag, size_t N, int direction) {
  const float PI = 3.14159265358979323846f;

  // Butterfly operations
  for (size_t stage_size = 2; stage_size <= N; stage_size *= 2) {
    const size_t half_stage = stage_size / 2;
    const float angle = direction * 2.0f * PI / stage_size;

    // Precompute twiddle factors for this stage
    float wr = 1.0f;
    float wi = 0.0f;
    const float wr_delta = cosf(angle);
    const float wi_delta = sinf(angle);

    for (size_t k = 0; k < half_stage; k++) {
      // Apply butterfly to all groups at this stage
      for (size_t group_start = k; group_start < N; group_start += stage_size) {
        const size_t i = group_start;
        const size_t j = group_start + half_stage;

        // Complex multiplication: (real[j] + i*imag[j]) * (wr + i*wi)
        const float temp_real = real[j] * wr - imag[j] * wi;
        const float temp_imag = real[j] * wi + imag[j] * wr;

        // Butterfly operation
        real[j] = real[i] - temp_real;
        imag[j] = imag[i] - temp_imag;
        real[i] = real[i] + temp_real;
        imag[i] = imag[i] + temp_imag;
      }

      // Update twiddle factor for next k (rotation)
      const float wr_old = wr;
      wr = wr_old * wr_delta - wi * wi_delta;
      wi = wr_old * wi_delta + wi * wr_delta;
    }
  }
}

void fft_forward(float* real, float* imag, size_t N) {
  bit_reverse_permute(real, imag, N);
  fft_radix2(real, imag, N, +1);
}

void fft_inverse(float* real, float* imag, size_t N) {
  bit_reverse_permute(real, imag, N);
  fft_radix2(real, imag, N, -1);

  // Scale by 1/N
  const float scale = 1.0f / N;
  for (size_t i = 0; i < N; i++) {
    real[i] *= scale;
    imag[i] *= scale;
  }
}

// DCT-II via FFT using reordering method
// Reference: Numerical Recipes Chapter 12.3, Press et al.
void dct_fft(const float* input, float* output, size_t N) {
  const float PI = 3.14159265358979323846f;

  // Allocate temporary arrays for N-point FFT
  float* real = new float[N];
  float* imag = new float[N];

  // Reorder input: even indices first, then odd indices reversed
  // [x[0], x[2], x[4], ...] followed by [x[N-1], x[N-3], x[N-5], ...]
  for (size_t i = 0; i < N / 2; i++) {
    real[i] = input[2 * i];           // Even indices: 0, 2, 4, ...
    real[N - 1 - i] = input[2 * i + 1]; // Odd indices reversed: N-1, N-3, ...
  }
  memset(imag, 0, N * sizeof(float));

  // Apply N-point FFT
  fft_forward(real, imag, N);

  // Extract DCT coefficients with phase correction
  // DCT[k] = Re{FFT[k] * exp(-j*pi*k/(2*N))} * normalization
  for (size_t k = 0; k < N; k++) {
    const float angle = -PI * k / (2.0f * N);
    const float wr = cosf(angle);
    const float wi = sinf(angle);

    // Complex multiplication: (real[k] + j*imag[k]) * (wr + j*wi)
    // Real part: real*wr - imag*wi
    const float dct_value = real[k] * wr - imag[k] * wi;

    // Apply DCT-II normalization
    if (k == 0) {
      output[k] = dct_value * sqrtf(1.0f / N);
    } else {
      output[k] = dct_value * sqrtf(2.0f / N);
    }
  }

  delete[] real;
  delete[] imag;
}

// IDCT (DCT-III) via FFT - inverse of the DCT-II reordering method
// Reference: Numerical Recipes Chapter 12.3, Press et al.
void idct_fft(const float* input, float* output, size_t N) {
  const float PI = 3.14159265358979323846f;

  // Allocate temporary arrays for N-point FFT
  float* real = new float[N];
  float* imag = new float[N];

  // Prepare FFT input with inverse phase correction
  // FFT[k] = DCT[k] * exp(+j*pi*k/(2*N)) / normalization
  // Note: DCT-III (inverse of DCT-II) requires factor of 2 for AC terms
  for (size_t k = 0; k < N; k++) {
    const float angle = PI * k / (2.0f * N);  // Positive angle for inverse
    const float wr = cosf(angle);
    const float wi = sinf(angle);

    // Inverse of DCT-II normalization with correct DCT-III scaling
    float scaled;
    if (k == 0) {
      scaled = input[k] / sqrtf(1.0f / N);
    } else {
      // DCT-III needs factor of 2 for AC terms
      scaled = input[k] / sqrtf(2.0f / N) * 2.0f;
    }

    // Complex multiplication: scaled * (wr + j*wi)
    real[k] = scaled * wr;
    imag[k] = scaled * wi;
  }

  // Apply inverse FFT
  fft_inverse(real, imag, N);

  // Unpack: reverse the reordering from DCT
  // Even output indices come from first half of FFT output
  // Odd output indices come from second half (reversed)
  for (size_t i = 0; i < N / 2; i++) {
    output[2 * i] = real[i];              // Even positions
    output[2 * i + 1] = real[N - 1 - i];  // Odd positions (reversed)
  }

  delete[] real;
  delete[] imag;
}