// Signal analysis for display: FFT, spectrogram, spectra, mode-energy bands.

export function fft(re: Float64Array, im: Float64Array) {
  const n = re.length;
  for (let i = 1, j = 0; i < n; i++) {
    let bit = n >> 1;
    for (; j & bit; bit >>= 1) j ^= bit;
    j ^= bit;
    if (i < j) {
      [re[i], re[j]] = [re[j], re[i]];
      [im[i], im[j]] = [im[j], im[i]];
    }
  }
  for (let len = 2; len <= n; len <<= 1) {
    const ang = (-2 * Math.PI) / len;
    const wr = Math.cos(ang);
    const wi = Math.sin(ang);
    for (let i = 0; i < n; i += len) {
      let cr = 1;
      let ci = 0;
      for (let k = 0; k < len / 2; k++) {
        const a = i + k;
        const b = a + len / 2;
        const tr = re[b] * cr - im[b] * ci;
        const ti = re[b] * ci + im[b] * cr;
        re[b] = re[a] - tr;
        im[b] = im[a] - ti;
        re[a] += tr;
        im[a] += ti;
        const nr = cr * wr - ci * wi;
        ci = cr * wi + ci * wr;
        cr = nr;
      }
    }
  }
}

export interface Spectrogram {
  nFrames: number;
  nBins: number;
  hop: number;
  win: number;
  db: Float32Array; // frame-major, nFrames x nBins
  maxDb: number;
}

export function spectrogram(x: Float32Array, win = 2048, hop = 256): Spectrogram {
  const nFrames = Math.max(1, Math.floor((x.length - win) / hop) + 1);
  const nBins = win / 2;
  const db = new Float32Array(nFrames * nBins);
  const re = new Float64Array(win);
  const im = new Float64Array(win);
  const w = new Float64Array(win);
  for (let i = 0; i < win; i++) w[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / win);
  let maxDb = -Infinity;
  for (let f = 0; f < nFrames; f++) {
    for (let i = 0; i < win; i++) {
      re[i] = (x[f * hop + i] ?? 0) * w[i];
      im[i] = 0;
    }
    fft(re, im);
    for (let b = 0; b < nBins; b++) {
      const v = 10 * Math.log10(re[b] * re[b] + im[b] * im[b] + 1e-300);
      db[f * nBins + b] = v;
      if (v > maxDb) maxDb = v;
    }
  }
  return { nFrames, nBins, hop, win, db, maxDb };
}

// Power spectrum (dB) of x[start, start+len) with a Hann window, zero padded to a power of two.
export function spectrum(x: Float32Array, start: number, len: number): { db: Float64Array; df: number; nfft: number } {
  len = Math.max(2, Math.min(len, x.length - start));
  let nfft = 1;
  while (nfft < 2 * len) nfft <<= 1;
  const re = new Float64Array(nfft);
  const im = new Float64Array(nfft);
  for (let i = 0; i < len; i++) re[i] = x[start + i] * (0.5 - 0.5 * Math.cos((2 * Math.PI * i) / (len - 1)));
  fft(re, im);
  const db = new Float64Array(nfft / 2);
  for (let b = 0; b < nfft / 2; b++) db[b] = 10 * Math.log10(re[b] * re[b] + im[b] * im[b] + 1e-300);
  return { db, df: 1 / nfft, nfft };
}

// Sum mode energies into log-spaced frequency bands.
export function modeBands(
  freqs: Float64Array,
  modeEnergy: Float32Array,
  nFrames: number,
  nBands = 48
): { edges: Float64Array; e: Float32Array; nBands: number } {
  const nModes = freqs.length;
  const fmin = Math.max(freqs[0], 1e-3) * 0.97;
  const fmax = freqs[nModes - 1] * 1.03;
  const edges = new Float64Array(nBands + 1);
  for (let b = 0; b <= nBands; b++) edges[b] = fmin * Math.pow(fmax / fmin, b / nBands);
  const band = new Int32Array(nModes);
  for (let k = 0; k < nModes; k++) {
    band[k] = Math.min(nBands - 1, Math.max(0, Math.floor((Math.log(freqs[k] / fmin) / Math.log(fmax / fmin)) * nBands)));
  }
  const e = new Float32Array(nFrames * nBands);
  for (let f = 0; f < nFrames; f++) {
    for (let k = 0; k < nModes; k++) e[f * nBands + band[k]] += modeEnergy[f * nModes + k];
  }
  return { edges, e, nBands };
}

// Relative RMS difference between two signals over [0, end).
export function relDiff(a: Float32Array, b: Float32Array, end: number): number {
  let num = 0;
  let den = 0;
  for (let i = 0; i < Math.min(end, a.length, b.length); i++) {
    num += (a[i] - b[i]) ** 2;
    den += a[i] ** 2;
  }
  return den > 0 ? Math.sqrt(num / den) : 0;
}

export function wavBlob(x: Float32Array, fs: number): Blob {
  let peak = 0;
  for (const v of x) peak = Math.max(peak, Math.abs(v));
  const g = peak > 0 ? 0.9 / peak : 1;
  const buf = new ArrayBuffer(44 + 2 * x.length);
  const dv = new DataView(buf);
  const w = (o: number, s: string) => [...s].forEach((c, i) => dv.setUint8(o + i, c.charCodeAt(0)));
  w(0, "RIFF");
  dv.setUint32(4, 36 + 2 * x.length, true);
  w(8, "WAVE");
  w(12, "fmt ");
  dv.setUint32(16, 16, true);
  dv.setUint16(20, 1, true);
  dv.setUint16(22, 1, true);
  dv.setUint32(24, fs, true);
  dv.setUint32(28, 2 * fs, true);
  dv.setUint16(32, 2, true);
  dv.setUint16(34, 16, true);
  w(36, "data");
  dv.setUint32(40, 2 * x.length, true);
  for (let i = 0; i < x.length; i++) dv.setInt16(44 + 2 * i, Math.round(32767 * Math.max(-1, Math.min(1, g * x[i]))), true);
  return new Blob([buf], { type: "audio/wav" });
}