import { test, expect } from "@playwright/test";
import { existsSync } from "node:fs";
import { romPath, symbolsPath, loadSymbols, openReference } from "./reference-fixture";

test.skip(!existsSync(romPath) || !existsSync(symbolsPath), "Needs private ROM and reference symbols.");

/** The synthesizer against the reference emulator's own sound output: the
 * cartridge's UpdateSound drives binjgb's APU once per frame, the port's
 * engine drives apu.ts with the same songs, and the two recordings are
 * compared window by window by their spectra (pitches and loudness), which
 * the consoles' different output filters leave alone. */
test("apu.ts sounds like the reference emulator's APU for songs, effects and cries", async ({ page }) => {
  test.setTimeout(300_000);
  const { at } = loadSymbols();
  const labels = {
    options: at("wOptions"),
    initSound: at("InitSound"),
    updateSound: at("UpdateSound"),
    playMusic: at("PlayMusic"),
    playSFX: at("PlaySFX"),
    playCry: at("PlayCry"),
  };
  await openReference(page);
  const results = await page.evaluate(async (labels) => {
    const { BinjgbEmulator } = await import("/src/emulator/binjgb.ts");
    const { AudioEngine } = await import("/src/game/audio/engine.ts");
    const { Apu, FRAME_CYCLES } = await import("/src/game/audio/apu.ts");
    const { RomReader } = await import("/src/game/rom/reader.ts");
    const rom = new Uint8Array(
      await (document.querySelector("#reference-fixture") as HTMLInputElement).files![0].arrayBuffer(),
    );
    const reference = new BinjgbEmulator();
    await reference.load(rom, document.createElement("canvas"));
    const inner = reference as unknown as {
      module: { HEAPU8: Uint8Array } & Record<string, (...args: number[]) => number>;
      handle: number;
      audio?: AudioContext;
    };
    const m = inner.module,
      h = inner.handle;
    const rate = inner.audio?.sampleRate ?? 44100;
    const write = (p: number, v: number) => reference.writeMemory(p, v & 255);
    const read = (p: number) => reference.readMemory(p);
    // Recorded reference sound, mixed to mono.
    let recorded: number[] = [];
    const run = (cycles: number) => {
      const target = m._emulator_get_ticks_f64(h) + cycles;
      for (let n = 0; n < 1000; n++) {
        const event = m._emulator_run_until_f64(h, target);
        if (event & 2) {
          const pointer = m._get_audio_buffer_ptr(h);
          const bytes = m.HEAPU8.subarray(pointer, pointer + 2048 * 2);
          for (let i = 0; i < 2048; i++) recorded.push((bytes[i * 2] + bytes[i * 2 + 1]) / 510 - 0.5);
        }
        if (event & 4) return;
      }
      throw Error("Reference cycle budget exceeded.");
    };
    // DI; SP; ld de, nn; call; sentinel; then spin. Returns the cycles the
    // call took, in 256-cycle steps.
    const call = (address: number, de = 0) => {
      const program = [0xf3, 0x31, 0xfe, 0xc7, 0x11, de & 255, de >> 8, 0xcd, address & 255, address >> 8, 0x3e, 0x42, 0xea, 0xf0, 0xc0, 0x18, 0xfe];
      write(0xc0f0, 0);
      program.forEach((v, i) => write(0xc000 + i, v));
      reference.setProgramCounter(0xc000);
      let cycles = 0;
      for (; cycles < 65536 && read(0xc0f0) !== 0x42; cycles += 256) run(256);
      if (read(0xc0f0) !== 0x42) throw Error(`Routine $${address.toString(16)} did not return`);
      return cycles;
    };
    const fft = (re: Float64Array) => {
      const n = re.length,
        im = new Float64Array(n);
      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]];
      }
      for (let len = 2; len <= n; len <<= 1) {
        const angle = (-2 * Math.PI) / len;
        for (let i = 0; i < n; i += len)
          for (let k = 0; k < len / 2; k++) {
            const wr = Math.cos(angle * k),
              wi = Math.sin(angle * 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;
          }
      }
      // Magnitudes up to 4 kHz, where the music's fundamentals and most
      // harmonics are.
      const bins = Math.floor((4000 / rate) * n);
      const out = new Float64Array(bins);
      for (let i = 1; i < bins; i++) out[i] = Math.hypot(re[i], im[i]);
      return out;
    };
    const SIZE = 4096;
    const spectra = (samples: ArrayLike<number>, start: number, count: number) => {
      const result: Float64Array[] = [];
      for (let w = 0; w < count; w++) {
        const frame = new Float64Array(SIZE);
        let mean = 0;
        for (let i = 0; i < SIZE; i++) mean += samples[start + w * SIZE + i] ?? 0;
        mean /= SIZE;
        for (let i = 0; i < SIZE; i++)
          frame[i] = ((samples[start + w * SIZE + i] ?? 0) - mean) * (0.5 - 0.5 * Math.cos((2 * Math.PI * i) / SIZE));
        result.push(fft(frame));
      }
      return result;
    };
    const cosine = (a: Float64Array, b: Float64Array) => {
      let dot = 0,
        na = 0,
        nb = 0;
      for (let i = 0; i < a.length; i++) {
        dot += a[i] * b[i];
        na += a[i] * a[i];
        nb += b[i] * b[i];
      }
      return na && nb ? dot / Math.sqrt(na * nb) : na === nb ? 1 : 0;
    };
    const cases = [
      { name: "Route 29", routine: labels.playMusic, de: 0x34, seconds: 8 },
      { name: "Johto wild battle", routine: labels.playMusic, de: 0x29, seconds: 8 },
      { name: "Pokémon Center", routine: labels.playMusic, de: 0x09, seconds: 6 },
      { name: "Goldenrod City", routine: labels.playMusic, de: 0x3d, seconds: 6 },
      { name: "Dragon's Den (noise, vibrato)", routine: labels.playMusic, de: 0x49, seconds: 6 },
      { name: "SFX_SAVE", routine: labels.playSFX, de: 0x25, seconds: 2 },
      { name: "Pikachu's cry", routine: labels.playCry, de: 24, seconds: 1.5 },
      { name: "Lugia's cry", routine: labels.playCry, de: 248, seconds: 2 },
    ];
    const results = [];
    for (const testCase of cases) {
      const engine = new AudioEngine(new RomReader(rom));
      const apu = new Apu(rate);
      write(labels.options, 0);
      call(labels.initSound);
      engine.initSound();
      call(testCase.routine, testCase.de);
      if (testCase.routine === labels.playMusic) engine.playMusic(testCase.de);
      else if (testCase.routine === labels.playSFX) engine.playSFX(testCase.de);
      else engine.playCry(testCase.de + 1);
      // Start both recordings at the first UpdateSound.
      recorded = [];
      const frames = Math.round(testCase.seconds * 59.7275);
      const native = new Float32Array(Math.ceil(((frames + 1) * FRAME_CYCLES * rate) / 4194304) + 4096);
      const left = new Float32Array(native.length),
        right = new Float32Array(native.length);
      let rendered = 0,
        due = 0;
      for (let frame = 0; frame < frames; frame++) {
        // One UpdateSound per video frame, as VBlank calls it.
        run(FRAME_CYCLES - call(labels.updateSound));
        engine.updateSound();
        for (const value of engine.drain()) apu.write(value >> 8, value & 0xff);
        due += (FRAME_CYCLES * rate) / 4194304;
        const count = Math.floor(due) - rendered;
        apu.render(left, right, rendered, count);
        rendered += count;
      }
      for (let i = 0; i < rendered; i++) native[i] = (left[i] + right[i]) / 2;
      // binjgb's buffer runs ahead by up to one 2048-sample chunk; align by
      // the best match of the loudness envelopes within that range.
      const windows = Math.floor(Math.min(rendered, recorded.length - 2048) / SIZE) - 1;
      let best = { score: -1, shift: 0 };
      for (let shift = -2048; shift <= 2048; shift += 64) {
        const a = spectra(native, Math.max(0, shift), Math.min(windows, 3)),
          b = spectra(recorded, Math.max(0, -shift), Math.min(windows, 3));
        const score = a.reduce((sum, spectrum, i) => sum + cosine(spectrum, b[i]), 0);
        if (score > best.score) best = { score, shift };
      }
      const a = spectra(native, Math.max(0, best.shift), windows),
        b = spectra(recorded, Math.max(0, -best.shift), windows);
      const scores = a.map((spectrum, i) => cosine(spectrum, b[i])).sort((x, y) => x - y);
      const rms = (x: ArrayLike<number>, n: number) => {
        let sum = 0;
        for (let i = 0; i < n; i++) sum += x[i] * x[i];
        return Math.sqrt(sum / Math.max(1, n));
      };
      results.push({
        name: testCase.name,
        nativeRms: rms(native, rendered),
        referenceRms: rms(recorded, recorded.length),
        recorded: recorded.length,
        rendered,
        windows,
        shift: best.shift,
        median: scores[Math.floor(scores.length / 2)],
        worst: scores[0],
      });
    }
    reference.dispose();
    return results;
  }, labels);
  console.log(JSON.stringify(results, null, 1));
  for (const result of results) {
    expect(result.windows, result.name).toBeGreaterThan(0);
    expect(result.median, result.name).toBeGreaterThan(0.85);
  }
});