import { test, expect } from "bun:test";
import { fft, stft, istft } from "../../src/dsp/fft.ts";
import { createRng } from "../../src/core/rng.ts";

test("fft of an impulse has all-ones magnitude", () => {
  const n = 1024;
  const re = new Float32Array(n);
  const im = new Float32Array(n);
  re[0] = 1;
  fft(re, im);
  for (let k = 0; k < n; k++) expect(Math.hypot(re[k]!, im[k]!)).toBeCloseTo(1, 6);
});

test("fft matches a direct DFT on a random signal", () => {
  const n = 64;
  const rng = createRng(3);
  const x = Array.from({ length: n }, () => rng.next() * 2 - 1);
  const re = Float32Array.from(x);
  const im = new Float32Array(n);
  fft(re, im);
  for (let k = 0; k < n; k++) {
    let sr = 0, si = 0;
    for (let t = 0; t < n; t++) {
      sr += x[t]! * Math.cos((-2 * Math.PI * k * t) / n);
      si += x[t]! * Math.sin((-2 * Math.PI * k * t) / n);
    }
    expect(re[k]!).toBeCloseTo(sr, 4);
    expect(im[k]!).toBeCloseTo(si, 4);
  }
});

test("stft follows torch semantics: frame count, bins, normalized=True scaling", () => {
  const nFft = 4096, hop = 1024, len = 50_000;
  const x = new Float32Array(len).fill(1); // DC: bin 0 of an interior frame = sum(hann)/sqrt(nFft) = nFft/2/sqrt(nFft)
  const s = stft(x, nFft, hop);
  expect(s.re.length).toBe(1 + Math.floor(len / hop));
  expect(s.re[0]!.length).toBe(nFft / 2 + 1);
  expect(s.re[10]![0]!).toBeCloseTo(nFft / 2 / Math.sqrt(nFft), 3);
  expect(Math.abs(s.re[10]![1]!)).toBeCloseTo(nFft / 4 / Math.sqrt(nFft), 3);
});

test("istft(stft(x)) reconstructs x (max err < 1e-4)", () => {
  const nFft = 4096, hop = 1024, len = 50_000;
  const rng = createRng(7);
  const x = new Float32Array(len);
  for (let i = 0; i < len; i++) x[i] = rng.next() * 2 - 1;
  const y = istft(stft(x, nFft, hop), nFft, hop, len);
  expect(y.length).toBe(len);
  let maxErr = 0;
  for (let i = 0; i < len; i++) maxErr = Math.max(maxErr, Math.abs(y[i]! - x[i]!));
  expect(maxErr).toBeLessThan(1e-4);
});