import { test, expect, beforeAll } from "bun:test";
import { join } from "node:path";
import { buildSong, loadSongDir } from "../../src/bench/dataset.ts";
import { createSeparator, STEMS, type Stem } from "../../src/pipeline/separate.ts";
import type { Audio } from "../../src/core/types.ts";

const MODEL = join(import.meta.dir, "../../models/demucs/htdemucs_6s.onnx");
const SLOW = 300_000;

let separate: (mix: Audio) => Promise<Record<Stem, Audio>>;
let data: Awaited<ReturnType<typeof loadSongDir>>;

beforeAll(async () => {
  const dir = await buildSong("t1b", 1);
  data = await loadSongDir(dir);
  separate = await createSeparator(new Uint8Array(await Bun.file(MODEL).arrayBuffer()), { intraOpNumThreads: 8 });
}, SLOW);

/** Sum of oracle stems for every track of the given role. */
function oracle(role: string): Audio {
  const len = data.mix.channels[0]!.length;
  const ch = [new Float32Array(len), new Float32Array(len)];
  for (const t of data.song.tracks) {
    if (t.role !== role) continue;
    const s = data.stems.get(t.name)!;
    for (let c = 0; c < 2; c++) for (let i = 0; i < len; i++) ch[c]![i]! += s.channels[c]![i]!;
  }
  return { sampleRate: data.mix.sampleRate, channels: ch };
}

/** SDR in dB over all channels: 10 log10(|ref|^2 / |ref - est|^2). */
function sdr(ref: Audio, est: Audio): number {
  let num = 0, den = 0;
  for (let c = 0; c < ref.channels.length; c++) {
    const r = ref.channels[c]!, e = est.channels[c]!;
    for (let i = 0; i < r.length; i++) {
      num += r[i]! * r[i]!;
      den += (r[i]! - e[i]!) ** 2;
    }
  }
  return 10 * Math.log10(num / Math.max(den, 1e-20));
}

let out: Record<Stem, Audio>;
let secsPerMin = 0;
const sdrs: Record<string, number> = {};

function energy(a: Audio): number {
  let e = 0;
  for (const ch of a.channels) for (const v of ch) e += v * v;
  return e;
}

/** Every sample finite, and the stems together carry a real share of the input's energy (not silence). */
function expectFiniteNonSilent(input: Audio, stems: Record<Stem, Audio>): void {
  let total = 0;
  for (const s of STEMS) {
    for (const ch of stems[s].channels) for (const v of ch) if (!Number.isFinite(v)) throw new Error(`${s}: non-finite sample ${v}`);
    total += energy(stems[s]);
  }
  expect(energy(input)).toBeGreaterThan(0);
  expect(total).toBeGreaterThan(0.1 * energy(input));
}

function sum(...as: Audio[]): Audio {
  const len = as[0]!.channels[0]!.length;
  return { sampleRate: 44100, channels: [0, 1].map((c) => {
    const o = new Float32Array(len);
    for (const a of as) for (let i = 0; i < len; i++) o[i]! += a.channels[c]![i]!;
    return o;
  }) };
}

beforeAll(async () => {
  const t0 = performance.now();
  out = await separate(data.mix);
  const minutes = data.mix.channels[0]!.length / data.mix.sampleRate / 60;
  secsPerMin = (performance.now() - t0) / 1000 / minutes;
  for (const [stem, role] of [["guitar", "guitar"], ["bass", "bass"], ["drums", "drums"]] as const) {
    sdrs[stem] = sdr(oracle(role), out[stem]);
    sdrs[`${stem}_mixBaseline`] = sdr(oracle(role), data.mix);
  }
  sdrs.guitarPlusPiano = sdr(oracle("guitar"), sum(out.guitar, out.piano));
  sdrs.guitarRefVsPianoStem = sdr(oracle("guitar"), out.piano);
  sdrs.guitarStemEnergy = energy(out.guitar);
  sdrs.pianoStemEnergy = energy(out.piano);
  console.log(`separation t1b/1: ${secsPerMin.toFixed(1)} s per minute of audio`, sdrs);
}, SLOW);

test("separates t1b/1: every stem stereo 44.1 kHz with exactly the input length", () => {
  const len = data.mix.channels[0]!.length;
  for (const s of STEMS) {
    expect(out[s].sampleRate).toBe(44100);
    expect(out[s].channels.length).toBe(2);
    for (const ch of out[s].channels) expect(ch.length).toBe(len);
  }
});

test("separates t1b/1: drums > 3 dB, bass beats the mix, guitar(+piano) > 3 dB and beats the mix", () => {
  // Measured (2026-10-30): drums 5.4 dB (mix -13.7), bass 1.67 dB (mix -2.8), guitar 1.85 dB (mix 1.87),
  // guitar+piano 4.7 dB. HTDemucs-6s routes ~30% of the GM steel guitar into its piano stem.
  expect(sdrs.drums!).toBeGreaterThan(3);
  expect(sdrs.drums!).toBeGreaterThan(sdrs.drums_mixBaseline!);
  expect(sdrs.bass!).toBeGreaterThan(sdrs.bass_mixBaseline!);
  expect(sdrs.bass!).toBeGreaterThan(0.5); // absolute floor: a silent estimate scores 0 dB
  expect(sdrs.guitarPlusPiano!).toBeGreaterThan(3);
  expect(sdrs.guitarPlusPiano!).toBeGreaterThan(sdrs.guitar_mixBaseline!);
});

test("separates t1b/1: guitar and piano stems are not swapped (guitar stem carries more energy)", () => {
  // Pins the manifest stem order (… guitar, piano). Energy, not SDR: the piano stem's SDR against the oracle
  // guitar (2.55 dB measured) exceeds the guitar stem's (1.85 dB) because both hold parts of the same guitar.
  // Measured energies: guitar 19521, piano 12721.
  expect(sdrs.guitarStemEnergy!).toBeGreaterThan(sdrs.pianoStemEnergy!);
});

// Brief's literal criterion, relaxed by controller ruling R7: HTDemucs-6s routes part of the SONiVOX GM guitar
// into its piano stem (synthetic-timbre gap), so guitar is 1.85 dB and bass 1.67 dB on t1b/1. Kept as an expected
// failure so it flips loudly if separation improves.
test.failing("brief criterion: guitar and bass stems each > 3 dB and > mix baseline", () => {
  expect(sdrs.guitar!).toBeGreaterThan(3);
  expect(sdrs.guitar!).toBeGreaterThan(sdrs.guitar_mixBaseline!);
  expect(sdrs.bass!).toBeGreaterThan(3);
  expect(sdrs.bass!).toBeGreaterThan(sdrs.bass_mixBaseline!);
});

test("a 3-second slice (shorter than one segment) returns 3 s stems", async () => {
  const n = 3 * 44100;
  const slice: Audio = { sampleRate: 44100, channels: data.mix.channels.map((c) => c.slice(44100, 44100 + n)) };
  const out = await separate(slice);
  for (const s of STEMS) for (const ch of out[s].channels) expect(ch.length).toBe(n);
  expectFiniteNonSilent(slice, out);
}, SLOW);

test("mono input is accepted and returns stereo stems of the same length", async () => {
  const n = 44100;
  const mono: Audio = { sampleRate: 44100, channels: [data.mix.channels[0]!.slice(0, n)] };
  const out = await separate(mono);
  for (const s of STEMS) {
    expect(out[s].channels.length).toBe(2);
    expect(out[s].channels[0]!.length).toBe(n);
  }
  expectFiniteNonSilent(mono, out);
}, SLOW);

test("rejects sample rates other than 44100", async () => {
  await expect(separate({ sampleRate: 22050, channels: [new Float32Array(100)] })).rejects.toThrow(/44100/);
});