// Find a short reference SFX inside a longer recording that has OTHER audio mixed
// over it (here: a clean music-off Poke Ball throw, against the longplay where the
// battle music is playing under it).
//
// Raw-waveform correlation is the wrong tool for that - the music is additive and the
// two encodes differ. What survives is the SFX's ENERGY SHAPE in a band where it
// dominates. So: per-band log-energy envelopes at a 2ms hop, mean-removed, then
// normalised cross-correlation. Reports the peak lag with a z-score against all lags,
// which is the only thing that says whether the match is real or the best of a lot of noise.
//
//   node xcorr-sfx.mjs <haystack.wav> <needle.wav> <haystackStartSec> [loBandHz] [hiBandHz]
import { readFileSync } from "node:fs";

function readWav(file) {
  const b = readFileSync(file);
  let p = 12, fmt = null, data = null;
  while (p + 8 <= b.length) {
    const id = b.toString("latin1", p, p + 4), sz = b.readUInt32LE(p + 4), body = p + 8;
    if (id === "fmt ") fmt = { ch: b.readUInt16LE(body + 2), sr: b.readUInt32LE(body + 4), bits: b.readUInt16LE(body + 14) };
    if (id === "data") { data = b.subarray(body, body + sz); break; }
    p = body + sz + (sz & 1);
  }
  const by = fmt.bits / 8, n = Math.floor(data.length / by / fmt.ch), x = new Float32Array(n);
  for (let i = 0; i < n; i += 1) {
    let s = 0;
    for (let c = 0; c < fmt.ch; c += 1) s += data.readInt16LE((i * fmt.ch + c) * by) / 32768;
    x[i] = s / fmt.ch;
  }
  return { x, sr: fmt.sr };
}

function fft(re, im) {
  const N = re.length;
  for (let i = 1, j = 0; i < N; i += 1) {
    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;
    for (let i = 0; i < N; i += len) for (let k = 0; k < len / 2; k += 1) {
      const wr = Math.cos(ang * k), wi = Math.sin(ang * k);
      const ur = re[i + k], ui = im[i + k];
      const vr = re[i + k + len / 2] * wr - im[i + k + len / 2] * wi;
      const vi = re[i + k + len / 2] * wi + im[i + k + len / 2] * wr;
      re[i + k] = ur + vr; im[i + k] = ui + vi;
      re[i + k + len / 2] = ur - vr; im[i + k + len / 2] = ui - vi;
    }
  }
}

// log-energy envelope in [lo,hi] Hz at a 2ms hop
function envelope(x, sr, lo, hi) {
  const N = 512, HOP = Math.round(0.002 * sr);
  const k0 = Math.max(1, Math.floor(lo / (sr / N))), k1 = Math.min(N / 2 - 1, Math.ceil(hi / (sr / N)));
  const out = [];
  for (let off = 0; off + N <= x.length; off += HOP) {
    const re = new Float64Array(N), im = new Float64Array(N);
    for (let i = 0; i < N; i += 1) re[i] = x[off + i] * (0.5 - 0.5 * Math.cos(2 * Math.PI * i / (N - 1)));
    fft(re, im);
    let e = 0;
    for (let k = k0; k <= k1; k += 1) e += re[k] * re[k] + im[k] * im[k];
    out.push(Math.log10(e + 1e-12));
  }
  return out;
}

const [hayF, needF, startArg, loArg, hiArg] = process.argv.slice(2);
const start = Number(startArg ?? 0), lo = Number(loArg ?? 1500), hi = Number(hiArg ?? 8000);
const hay = readWav(hayF), need = readWav(needF);
const HOP = 0.002;

const H = envelope(hay.x, hay.sr, lo, hi);
const Nv = envelope(need.x, need.sr, lo, hi);
const nm = Nv.reduce((a, b) => a + b, 0) / Nv.length;
const nz = Nv.map((v) => v - nm);
const nEn = Math.sqrt(nz.reduce((a, b) => a + b * b, 0));

const scores = [];
for (let off = 0; off + nz.length <= H.length; off += 1) {
  let hm = 0;
  for (let i = 0; i < nz.length; i += 1) hm += H[off + i];
  hm /= nz.length;
  let dot = 0, hEn = 0;
  for (let i = 0; i < nz.length; i += 1) { const h = H[off + i] - hm; dot += h * nz[i]; hEn += h * h; }
  scores.push([off * HOP, dot / (Math.sqrt(hEn) * nEn + 1e-12)]);
}

const vals = scores.map((s) => s[1]);
const mean = vals.reduce((a, b) => a + b, 0) / vals.length;
const sd = Math.sqrt(vals.reduce((a, b) => a + (b - mean) * (b - mean), 0) / vals.length);
const ranked = [...scores].sort((a, b) => b[1] - a[1]);

console.log(`band ${lo}-${hi}Hz, needle ${(need.x.length / need.sr).toFixed(3)}s, haystack ${(hay.x.length / hay.sr).toFixed(1)}s from ${start}s`);
console.log(`peak: haystack +${ranked[0][0].toFixed(3)}s  =>  SOURCE ${(start + ranked[0][0]).toFixed(3)}s   r=${ranked[0][1].toFixed(3)}  z=${((ranked[0][1] - mean) / sd).toFixed(1)}`);
console.log("top 6:");
for (const [t, r] of ranked.slice(0, 6)) console.log(`  src ${(start + t).toFixed(3)}s  r=${r.toFixed(3)}  z=${((r - mean) / sd).toFixed(1)}`);