// Tests for FFT-based DCT/IDCT implementation
// Verifies correctness against reference O(N²) implementation

#include "audio/fft.h"

#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstring>

// Reference O(N²) DCT-II implementation (from original code)
static void dct_reference(const float* input, float* output, size_t N) {
  const float PI = 3.14159265358979323846f;

  for (size_t k = 0; k < N; k++) {
    float sum = 0.0f;
    for (size_t n = 0; n < N; n++) {
      sum += input[n] * cosf((PI / N) * k * (n + 0.5f));
    }

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

// Reference O(N²) IDCT implementation (from original code)
static void idct_reference(const float* input, float* output, size_t N) {
  const float PI = 3.14159265358979323846f;

  for (size_t n = 0; n < N; ++n) {
    float sum = input[0] / 2.0f;
    for (size_t k = 1; k < N; ++k) {
      sum += input[k] * cosf((PI / N) * k * (n + 0.5f));
    }
    output[n] = sum * (2.0f / N);
  }
}

// Compare two arrays with tolerance
// Note: FFT-based DCT accumulates slightly more rounding error than O(N²) direct method
// A tolerance of 1e-3 is still excellent for audio applications (< -60 dB error)
static bool arrays_match(const float* a,
                         const float* b,
                         size_t N,
                         float tolerance = 1e-3f) {
  for (size_t i = 0; i < N; i++) {
    const float diff = fabsf(a[i] - b[i]);
    if (diff > tolerance) {
      fprintf(stderr,
              "Mismatch at index %zu: %.6f vs %.6f (diff=%.6e)\n",
              i,
              a[i],
              b[i],
              diff);
      return false;
    }
  }
  return true;
}

// Test 1: DCT correctness (FFT-based vs reference)
static void test_dct_correctness() {
  printf("Test 1: DCT correctness (FFT vs reference O(N²))...\n");

  const size_t N = 512;
  float input[N];
  float output_ref[N];
  float output_fft[N];

  // Test case 1: Impulse at index 0
  memset(input, 0, N * sizeof(float));
  input[0] = 1.0f;

  dct_reference(input, output_ref, N);
  dct_fft(input, output_fft, N);

  assert(arrays_match(output_ref, output_fft, N));
  printf("  ✓ Impulse test passed\n");

  // Test case 2: Impulse at middle
  memset(input, 0, N * sizeof(float));
  input[N / 2] = 1.0f;

  dct_reference(input, output_ref, N);
  dct_fft(input, output_fft, N);

  assert(arrays_match(output_ref, output_fft, N));
  printf("  ✓ Middle impulse test passed\n");

  // Test case 3: Sinusoidal input
  for (size_t i = 0; i < N; i++) {
    input[i] = sinf(2.0f * 3.14159265358979323846f * 5.0f * i / N);
  }

  dct_reference(input, output_ref, N);
  dct_fft(input, output_fft, N);

  assert(arrays_match(output_ref, output_fft, N));
  printf("  ✓ Sinusoidal input test passed\n");

  // Test case 4: Random-ish input (deterministic)
  for (size_t i = 0; i < N; i++) {
    input[i] = sinf(i * 0.1f) * cosf(i * 0.05f);
  }

  dct_reference(input, output_ref, N);
  dct_fft(input, output_fft, N);

  assert(arrays_match(output_ref, output_fft, N));
  printf("  ✓ Complex input test passed\n");

  printf("Test 1: PASSED ✓\n\n");
}

// Test 2: IDCT correctness (FFT-based vs reference)
static void test_idct_correctness() {
  printf("Test 2: IDCT correctness (FFT vs reference O(N²))...\n");

  const size_t N = 512;
  float input[N];
  float output_ref[N];
  float output_fft[N];

  // Test case 1: DC component only
  memset(input, 0, N * sizeof(float));
  input[0] = 1.0f;

  idct_reference(input, output_ref, N);
  idct_fft(input, output_fft, N);

  assert(arrays_match(output_ref, output_fft, N));
  printf("  ✓ DC component test passed\n");

  // Test case 2: Single frequency bin
  memset(input, 0, N * sizeof(float));
  input[10] = 1.0f;

  idct_reference(input, output_ref, N);
  idct_fft(input, output_fft, N);

  assert(arrays_match(output_ref, output_fft, N));
  printf("  ✓ Single bin test passed\n");

  // Test case 3: Mixed frequencies
  for (size_t i = 0; i < N; i++) {
    input[i] = (i % 10 == 0) ? 1.0f : 0.0f;
  }

  idct_reference(input, output_ref, N);
  idct_fft(input, output_fft, N);

  assert(arrays_match(output_ref, output_fft, N));
  printf("  ✓ Mixed frequencies test passed\n");

  printf("Test 2: PASSED ✓\n\n");
}

// Test 3: Round-trip (DCT → IDCT should recover original)
static void test_roundtrip() {
  printf("Test 3: Round-trip (DCT → IDCT = identity)...\n");

  const size_t N = 512;
  float input[N];
  float dct_output[N];
  float reconstructed[N];

  // Test case 1: Sinusoidal input
  for (size_t i = 0; i < N; i++) {
    input[i] = sinf(2.0f * 3.14159265358979323846f * 3.0f * i / N);
  }

  dct_fft(input, dct_output, N);
  idct_fft(dct_output, reconstructed, N);

  assert(arrays_match(input, reconstructed, N));
  printf("  ✓ Sinusoidal round-trip passed\n");

  // Test case 2: Complex signal
  for (size_t i = 0; i < N; i++) {
    input[i] = sinf(i * 0.1f) * cosf(i * 0.05f) + cosf(i * 0.03f);
  }

  dct_fft(input, dct_output, N);
  idct_fft(dct_output, reconstructed, N);

  assert(arrays_match(input, reconstructed, N));
  printf("  ✓ Complex signal round-trip passed\n");

  printf("Test 3: PASSED ✓\n\n");
}

// Test 4: Output known values for JavaScript comparison
static void test_known_values() {
  printf("Test 4: Known values (for JavaScript verification)...\n");

  const size_t N = 512;
  float input[N];
  float output[N];

  // Simple test case: impulse at index 0
  memset(input, 0, N * sizeof(float));
  input[0] = 1.0f;

  dct_fft(input, output, N);

  printf("  DCT of impulse at 0:\n");
  printf("    output[0] = %.8f (expected ~0.04419417)\n", output[0]);
  printf("    output[1] = %.8f (expected ~0.04419417)\n", output[1]);
  printf("    output[10] = %.8f (expected ~0.04419417)\n", output[10]);

  // IDCT test
  memset(input, 0, N * sizeof(float));
  input[0] = 1.0f;

  idct_fft(input, output, N);

  printf("  IDCT of DC component:\n");
  printf("    output[0] = %.8f (expected ~0.04419417)\n", output[0]);
  printf("    output[100] = %.8f (expected ~0.04419417)\n", output[100]);
  printf("    output[511] = %.8f (expected ~0.04419417)\n", output[511]);

  printf("Test 4: PASSED ✓\n");
  printf("(Copy these values to JavaScript test for verification)\n\n");
}

int main() {
  printf("===========================================\n");
  printf("FFT-based DCT/IDCT Test Suite\n");
  printf("===========================================\n\n");

  test_dct_correctness();
  test_idct_correctness();
  test_roundtrip();
  test_known_values();

  printf("===========================================\n");
  printf("All tests PASSED ✓\n");
  printf("===========================================\n");

  return 0;
}