STSheets/src/bench/real.ts

// Real-song benchmark: full pipeline on a team recording, scored per role against the GP-exported MIDI warped through the saved alignment.
import { join } from "node:path";
import { mkdir } from "node:fs/promises";
import type { Role, TimedNote } from "../core/types.ts";
import { createPipeline } from "../pipeline/run.ts";
import { matchNotes, type MatchResult } from "../eval/notes.ts";
import { DRUM_CLASSES, ROLE_ORDER, drumClass, toClassPitch, type DrumClass } from "../eval/score.ts";
import { readManifest, readAlignment, loadSong, trackKeys } from "../ref/manifest.ts";
import { makeWarp, type Anchor } from "../align/warp.ts";
import { autoAlign, autoSummary, constantOffsetBaseline, fitFromCues, recordingCues, type Baseline } from "../align/auto.ts";
import { fitSummary, type FitReport } from "../align/fit.ts";

export interface RealSongResult {
  id: string; title: string;
  /**
   * source: "saved" = reference/<id>/alignment.json, "computed" = auto-aligned this run (--auto). chosen: how the auto
   * alignment was made ("dtw" or the constant-offset "baseline"), null when unknown (saved from the page).
   * baseline: the constant-offset baseline computed this run, for comparison.
   */
  alignment: {
    source: "saved" | "computed"; manual: number; auto: number; savedAt: string | null;
    chosen: "baseline" | "dtw" | null; baseline: Baseline | { offsetSec: number; precision: number | null }; dtwPrecision: number | null;
  };
  fit: FitReport;
  perRole: Partial<Record<Role, MatchResult>>;
  /** F1 at the lenient +-100 ms tolerance, per role (null when the role has no notes on either side). */
  lenient: Partial<Record<Role, number | null>>;
  drumsPerClass: Record<DrumClass, MatchResult> | null;
  excludedTracks: string[];
  /** Capo per track from the manifest; informational only (not applied to scoring). */
  capo: Record<string, number>;
  seconds: number;
}

const MODELS = { demucs: "models/demucs/htdemucs_6s.onnx", basicPitch: "models/basic-pitch/nmp.onnx" };
const SF2 = join(import.meta.dir, "../../node_modules/@coderline/alphatab/dist/soundfont/sonivox.sf2");
const TOL = 0.05, LENIENT_TOL = 0.1;

export async function runReal(opts: { runId: string; auto: boolean; root?: string; out?: string; resultsDir?: string; repeatVote?: boolean }): Promise<RealSongResult[]> {
  const root = opts.root ?? "reference", out = opts.out ?? "runs", rdir = opts.resultsDir ?? "results";
  const songs = await readManifest(root);
  const pipeline = await createPipeline({
    demucs: new Uint8Array(await Bun.file(MODELS.demucs).arrayBuffer()),
    basicPitch: new Uint8Array(await Bun.file(MODELS.basicPitch).arrayBuffer()),
  }, { repeatVote: opts.repeatVote ?? true });
  const sf2 = opts.auto ? new Uint8Array(await Bun.file(SF2).arrayBuffer()) : null;
  const results: RealSongResult[] = [];
  const skipped: SkippedSong[] = [];
  const skip = (id: string, reason: string) => { skipped.push({ id, reason }); console.log(`${id}: skipped (${reason})`); };
  for (const entry of songs) {
    let saved: Awaited<ReturnType<typeof readAlignment>>;
    try {
      saved = await readAlignment(entry.id, root);
      if (saved) makeWarp(saved.anchors); // validates the anchors
    } catch (e) { skip(entry.id, (e as Error).message); continue; }
    if (!saved && !opts.auto) { skip(entry.id, "no alignment.json; pass --auto to compute one"); continue; }
    const t0 = performance.now();
    const { smf, audio } = await loadSong(entry.id, root);
    const keys = trackKeys(smf.tracks);
    const roleOf = (i: number): Role | null => entry.roles[keys[i]!] ?? null;
    const roles = Object.fromEntries(keys.map((k) => [k, entry.roles[k] ?? null]));

    const cues = recordingCues(audio);
    let anchors: Anchor[], alignment: RealSongResult["alignment"];
    if (saved) {
      anchors = saved.anchors;
      alignment = {
        source: "saved", manual: 0, auto: 0, savedAt: saved.savedAt ?? null,
        chosen: saved.auto?.chosen ?? null, baseline: constantOffsetBaseline(cues, smf, roles), dtwPrecision: saved.auto?.dtwPrecision ?? null,
      };
    } else {
      const r = autoAlign(audio, smf, roles, [], sf2!);
      anchors = r.anchors;
      alignment = { source: "computed", manual: 0, auto: 0, savedAt: null, chosen: r.chosen, baseline: r.baseline, dtwPrecision: r.dtwPrecision };
    }
    alignment.manual = anchors.filter((a) => a.kind === "manual").length;
    alignment.auto = anchors.filter((a) => a.kind === "auto").length;
    const fit = fitFromCues(cues, smf, roles, anchors);
    const warp = makeWarp(anchors);

    const res = await pipeline({ mix: audio }, { oracleStems: false, oracleTempo: false, oracleNotes: false });
    const refOf = (role: Role): TimedNote[] => smf.tracks.flatMap((t, i) => roleOf(i) === role
      ? t.notes.map((n) => ({ onset: warp.toAudio(n.onset), offset: warp.toAudio(n.offset), pitch: n.pitch })) : []);
    const estOf = (role: Role): TimedNote[] => res.output.tracks.filter((t) => t.role === role).flatMap((t) => t.notes);

    const perRole: RealSongResult["perRole"] = {}, lenient: RealSongResult["lenient"] = {};
    let drumsPerClass: RealSongResult["drumsPerClass"] = null;
    for (const role of ROLE_ORDER) {
      let r = refOf(role), e = estOf(role);
      if (role === "drums") {
        const rawR = r, rawE = e;
        r = r.map(toClassPitch); e = e.map(toClassPitch);
        if (rawR.length + rawE.length > 0) {
          drumsPerClass = {} as Record<DrumClass, MatchResult>;
          for (const c of Object.keys(DRUM_CLASSES) as DrumClass[]) {
            drumsPerClass[c] = matchNotes(rawR.filter((n) => drumClass(n.pitch) === c).map(toClassPitch), rawE.filter((n) => drumClass(n.pitch) === c).map(toClassPitch), { onsetTol: TOL });
          }
        }
      }
      perRole[role] = matchNotes(r, e, { onsetTol: TOL });
      lenient[role] = matchNotes(r, e, { onsetTol: LENIENT_TOL }).f1;
    }
    const seconds = (performance.now() - t0) / 1000;
    const result: RealSongResult = {
      id: entry.id, title: entry.title,
      alignment,
      fit, perRole, lenient, drumsPerClass,
      excludedTracks: smf.tracks.map((t, i) => (roleOf(i) ? null : keys[i]!)).filter((k): k is string => k !== null),
      capo: entry.capo ?? {}, seconds,
    };
    const d = join(out, `real-${opts.runId}`, entry.id);
    await mkdir(d, { recursive: true });
    await Bun.write(join(d, "notes.json"), JSON.stringify({ grid: res.output.grid, tracks: res.output.tracks, warnings: res.warnings, anchors, seconds }));
    await Bun.write(join(d, "out.gp"), res.gp);
    await Bun.write(join(d, "score.json"), JSON.stringify(result, null, 1));
    results.push(result);
    console.log(`${entry.id}: guitar F1 ${fmt(perRole.guitar?.f1)} (${seconds.toFixed(1)}s)`);
  }
  await mkdir(rdir, { recursive: true });
  await Bun.write(join(rdir, `real-${opts.runId}.json`), JSON.stringify(results, (k, v) => (k === "matches" ? undefined : v), 1));
  await Bun.write(join(rdir, `real-${opts.runId}.md`), renderRealReport(opts.runId, results, skipped));
  return results;
}

const NA = "n/a";
const fmt = (x: number | null | undefined) => (x == null ? NA : x.toFixed(2));
const row = (cells: string[]) => `| ${cells.join(" | ")} |`;

export interface SkippedSong { id: string; reason: string }

export function renderRealReport(runId: string, results: RealSongResult[], skipped: SkippedSong[] = []): string {
  const o: string[] = [`# STSheets real-song benchmark: ${runId}`, ""];
  o.push("> Reference = the team's Guitar Pro-exported MIDI warped onto the recording by the saved alignment; scores are note-onset matches (pitch must equal; drums compared by class). " +
    "The fit line reports the *consistency of the alignment*, not an independent validation of it: precision = share of the recording's onset-flux peaks with a warped MIDI onset within ±70 ms " +
    "(flag < 60% per 30 s window), lift = precision minus its value with the MIDI shifted off-grid (flag < 0.25), a symbolic-chroma check for whole-beat/bar shifts, " +
    "and \"MIDI onsets audible as attacks\" (information only; a full mix exposes far fewer attacks than the score has notes). " +
    "A poor alignment lowers every score below. Capo is informational and not applied.", "");
  for (const r of results) {
    const a = r.alignment;
    const counts = `${a.manual} manual / ${a.auto} auto anchors`;
    const src = a.source === "computed" ? `computed this run (${counts})` : `saved ${a.savedAt ?? "(time unknown)"} (${counts})`;
    const pc = (x: number | null) => (x === null ? NA : `${(x * 100).toFixed(1)}%`);
    const off = `${a.baseline.offsetSec >= 0 ? "+" : "−"}${Math.abs(a.baseline.offsetSec).toFixed(2)} s`;
    const method = a.chosen ? autoSummary({ chosen: a.chosen, baseline: a.baseline, dtwPrecision: a.dtwPrecision })
      : `method unknown (saved from the page) · constant-offset baseline ${off}: precision ${pc(a.baseline.precision)}`;
    o.push(`## ${r.title} (${r.id})`, "");
    o.push(`- Alignment: ${src}`, `- Method: ${method}`, `- Consistency of the alignment: ${fitSummary(r.fit)}`);
    const capos = Object.entries(r.capo);
    o.push(`- Capo: ${capos.length ? capos.map(([k, v]) => `${k}: ${v}`).join(", ") : "none listed"}`);
    o.push(`- Excluded tracks (no role): ${r.excludedTracks.length ? r.excludedTracks.join(", ") : "none"}`, `- Pipeline + scoring: ${r.seconds.toFixed(1)} s`, "");
    const head = ["Role", "Ref notes", "Est notes", "P ±50 ms", "R ±50 ms", "F1 ±50 ms", "F1 ±100 ms"];
    o.push(row(head), row(head.map(() => "---")));
    for (const role of ROLE_ORDER) {
      const m = r.perRole[role];
      if (!m) continue;
      o.push(row([role, String(m.tp + m.fn), String(m.tp + m.fp), fmt(m.precision), fmt(m.recall), fmt(m.f1), fmt(r.lenient[role])]));
    }
    o.push("");
    if (r.drumsPerClass) {
      const dh = ["Drum class", "Ref", "Est", "F1 ±50 ms"];
      o.push(row(dh), row(dh.map(() => "---")));
      for (const [c, m] of Object.entries(r.drumsPerClass)) o.push(row([c, String(m.tp + m.fn), String(m.tp + m.fp), fmt(m.f1)]));
      o.push("");
    }
  }
  if (skipped.length) o.push(`Skipped: ${skipped.map((s) => `${s.id} (${s.reason})`).join("; ")}`, "");
  return o.join("\n");
}