scripts/test_sc2_vendor.py

"""CPU acceptance for the vendored BarScript serving path.

Three questions, and nothing that needs a GPU or the network:

(a) does the BarScript model load from the release package ALONE -- with the
    development worktree stripped out of ``sys.path`` -- and does it read
    nothing outside that package?
(b) does it generate a chart whose TJA the vendored parser reads back?
(c) is the 1.x serving path the Space's seven published models use bit-for-bit
    unaffected: the library files untouched, the public signatures unchanged,
    the two namespaces isolated, and a real 1.x decode identical before and
    after ``softchart.sc2`` is imported?

Run:  python spaces/scripts/test_sc2_vendor.py
      pytest spaces/scripts/test_sc2_vendor.py

(c)'s decode check needs ``JacobLinCool/softchart-v17`` in the local Hub
cache; it reports SKIP, and does not fail, when the cache is cold, because it
never goes to the network.

The order-sensitive checks run in subprocesses of this file (``--case``): a
1.x baseline must be measured in a process where ``softchart.sc2`` was never
imported, and the package-only load must happen before anything else can put
the worktree back on the path.
"""

from __future__ import annotations

import argparse
import hashlib
import inspect
import json
import os
import subprocess
import sys
import warnings

SPACES_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
REPO_ROOT = os.path.dirname(SPACES_DIR)
SC2_DIR = os.path.join(SPACES_DIR, "softchart", "sc2")
LEGACY_REPO = os.environ.get("SC_LEGACY_TEST_REPO", "JacobLinCool/softchart-v17")

if SPACES_DIR not in sys.path:
    sys.path.insert(0, SPACES_DIR)


# ---------------------------------------------------------------------------
# (c) tripwires: the 1.x library as it stood when 2.0 was vendored in.
# Vendoring added files; it changed none. A mismatch here means someone edited
# a module the seven published checkpoints decode through, and their charts are
# no longer the baseline the user keeps them around to compare against.

LEGACY_MD5 = {
    "__init__.py": "d41d8cd98f00b204e9800998ecf8427e",
    "baseline.py": "875b63450c59bb62f5d50965ec7b218c",
    "data.py": "72e19832477dd49e2bb0dc731e17b3ac",
    "evaluate.py": "13840e66287b1113202d4ea0152f62ba",
    "fonts.py": "a6b1c5c1d3f9ab852fb551d49031d3a2",
    "generate.py": "b3ad2aa451ed413f1d9d64eafb072cd8",
    "grid.py": "570ae31b7120369781cd02c20de74ffd",
    "hf.py": "522c2257a402bd0d41727c6dee6732db",
    "model.py": "5e8ee5091629c9ff718d7c136d5f97bc",
    "preprocess.py": "1d765a119d909f2aa74d594f4ca6b9c0",
    "preview_audio.py": "b03aec51b160ccb33d8d553dd936d2c4",
    "proxies.py": "918c72c22d48e4ccb762e0fe814d0bfe",
    "rhythm.py": "11cca7c58809b339f7384dcc0bf1fcaa",
    "tja.py": "8e3905fc3b432026f034548ba8422b06",
    "tja_image.py": "afa678700e033206d9b80f9510dfc844",
    "train.py": "9993791112402dfee700e62cb860c470",
    "vocab.py": "46c4df35768ea7eee5cebccf63cc6c9d",
}

# The 1.x entry points app.py calls. Signature text, not just presence: a
# reordered or renamed keyword is a silent behaviour change for the app.
LEGACY_SIGNATURES = {
    "generate.generate_song": (
        "(model, mel, course, level=None, density_bucket=None, device='cuda', "
        "greedy=False, temperature=1.0, top_p=0.95, seed=0, batch_windows=8, "
        "cfg_w=0.0, use_ctx=None, lattice=False, style=None, sync_band=None, "
        "plan=None, on_progress=None)"),
    "generate.generate_song_slot": (
        "(model, mel, grid, course, level=None, density_bucket=None, "
        "device='cuda', greedy=False, temperature=1.0, top_p=0.95, seed=0, "
        "batch_windows=8, plan=None, on_progress=None)"),
    "generate.load_hf": "(repo_or_dir, device='cuda')",
    "tja.write_tja_slots": (
        "(gen, grid, title, course, level, wave, out_path=None, "
        "balloon_count=10, plan=None)"),
}

# 1.x constants the app reads, and the vocabulary geometry every published
# checkpoint's embedding table is indexed by.
LEGACY_CONSTANTS = {
    "SR": 22050, "N_FFT": 2048, "HOP": 256, "N_MELS": 128,
    "WINDOW": 1728, "MAX_TGT": 1536, "SLOTS": 96, "N_DENS": 16,
    "N_LEVELS": 12, "VOCAB_SIZE": 1810,
}

MEL_FRAMES = 1500  # ~17.4 s at 86.13 fps -- one decode window plus change


def _md5(path):
    h = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def _fake_mel(frames=MEL_FRAMES, seed=0):
    """A deterministic stand-in for a log-mel spectrogram.

    The acceptance here is mechanical -- loads, decodes, exports, re-parses --
    so the audio need only be well-formed and reproducible. Chart quality is
    the model card's subject, not this file's.
    """
    import numpy as np

    rng = np.random.default_rng(seed)
    t = np.arange(frames) / 86.1328125
    pulse = np.sin(2 * np.pi * 2.0 * t)[None, :] + 1.0
    return np.log(1e-5 + np.abs(rng.standard_normal((128, frames)) * 0.1
                                + pulse)).astype(np.float32)


def _run_case(case):
    """Run one order-sensitive case in a clean subprocess of this file."""
    env = dict(os.environ, HF_HUB_OFFLINE="1", PYTHONWARNINGS="ignore")
    proc = subprocess.run([sys.executable, os.path.abspath(__file__),
                           "--case", case],
                          capture_output=True, text=True, env=env,
                          cwd=REPO_ROOT)
    if proc.returncode != 0:
        raise AssertionError(
            f"case {case!r} failed (exit {proc.returncode}):\n"
            f"{proc.stdout[-4000:]}\n{proc.stderr[-4000:]}")
    for line in reversed(proc.stdout.splitlines()):
        if line.startswith("RESULT "):
            return json.loads(line[len("RESULT "):])
    raise AssertionError(f"case {case!r} printed no RESULT:\n{proc.stdout}")


# ---------------------------------------------------------------------------
# (a) package-only load


def case_package_only_load():
    """Load with the development worktree removed from the path, and watch every read.

    ``sys.path`` is scrubbed of anything under ``.codex/worktrees`` first, so
    an accidental ``import softchart.<module>`` resolving to the frozen
    shipping tree would fail rather than quietly succeed. Python-level file
    opens are recorded; every one that lands inside this repository must be
    inside the package directory or the vendored ``softchart/sc2`` tree.
    """
    import builtins

    # Put the development worktree on the path first, so the scrub has something
    # real to remove: a load that only passes because the worktree happened to
    # be absent would prove nothing.
    worktree_src = os.path.join(REPO_ROOT, ".codex", "worktrees", "v18-8h",
                                "src")
    if os.path.isdir(worktree_src):
        sys.path.insert(0, worktree_src)
    removed = [p for p in sys.path if ".codex" + os.sep + "worktrees" in p]
    sys.path[:] = [p for p in sys.path if p not in removed]
    for name in [n for n in sys.modules
                 if n == "softchart" or n.startswith("softchart.")]:
        del sys.modules[name]

    # modules that exist only in the worktree tree must now be unreachable
    try:
        __import__("softchart.barscript")
        worktree_reachable = True
    except ImportError:
        worktree_reachable = False
    for name in [n for n in sys.modules
                 if n == "softchart" or n.startswith("softchart.")]:
        del sys.modules[name]

    opened = []
    real_open = builtins.open

    def watched_open(file, *a, **kw):
        try:
            opened.append(os.path.abspath(os.fspath(file)))
        except TypeError:
            pass
        return real_open(file, *a, **kw)

    from softchart import sc2_loader as loader

    package_dir = os.path.abspath(loader.default_package_dir())
    builtins.open = watched_open
    try:
        model, config, spec = loader.load_release(package_dir, device="cpu")
    finally:
        builtins.open = real_open

    stray = [p for p in opened
             if p.startswith(REPO_ROOT + os.sep)
             and not p.startswith(package_dir + os.sep)
             and not p.startswith(SC2_DIR + os.sep)]
    outside = [m.__file__ for name, m in sys.modules.items()
               if name.startswith("softchart.sc2.")
               and getattr(m, "__file__", None)
               and not os.path.abspath(m.__file__).startswith(SC2_DIR + os.sep)]

    return {
        "package_dir": package_dir,
        "worktree_entries_removed": removed,
        "worktree_src_exists": os.path.isdir(worktree_src),
        "worktree_reachable_after_scrub": worktree_reachable,
        "release_name": config["release_name"],
        "parameters": int(config["weights"]["parameters"]),
        "vocab_size": int(config["vocab"]["size"]),
        "has_sync_forced": bool(getattr(model, "_has_sync", False)),
        "has_axis_knobs": bool(getattr(model, "_has_axis_knobs", False)),
        "motif_hard": config["train_provenance"]["motif_constraint"][
            "resolved_motif_hard"],
        "spec_labels_version": spec.labels_version,
        "files_opened_in_repo": len([p for p in opened
                                     if p.startswith(REPO_ROOT + os.sep)]),
        "stray_reads": stray,
        "sc2_modules_outside_vendor_dir": outside,
    }


def test_a_package_only_load():
    r = _run_case("package-only-load")
    if r["worktree_src_exists"]:
        assert r["worktree_entries_removed"], "the scrub removed nothing"
    assert r["worktree_reachable_after_scrub"] is False, (
        "the development worktree is still importable after the scrub; the load below "
        "proves nothing")
    assert not r["stray_reads"], f"read outside the package: {r['stray_reads']}"
    assert not r["sc2_modules_outside_vendor_dir"], r[
        "sc2_modules_outside_vendor_dir"]
    assert r["vocab_size"] == 2071, r["vocab_size"]
    # config.serving_prefix_fix: the recorded sync_token flag is wrong and the
    # loader must override it, or every sync bucket decodes identically.
    assert r["has_sync_forced"] is True
    assert r["has_axis_knobs"] is True
    assert r["motif_hard"] is False, "package resolved motif_hard off"
    assert r["spec_labels_version"] == "barscript_labels_v1"
    return r


# ---------------------------------------------------------------------------
# (b) generate and re-parse


def test_b_generate_and_reparse():
    from softchart import sc2_loader as loader
    from softchart.sc2.trace import chart_to_trace, parse_tja

    sc = loader.load(device="cpu")

    # the front end app.py computes must be the one the package was trained on
    from softchart.vocab import HOP, N_FFT, N_MELS, SR

    loader.check_preprocessing(sc.package_dir, sr=SR, n_fft=N_FFT,
                               hop_length=HOP, n_mels=N_MELS, fmin=20.0,
                               fmax=SR / 2)

    out = {}
    for course, level, dens in (("oni", 8, 6), ("ura", 10, 6), ("easy", 3, 1)):
        r = sc.generate(_fake_mel(), course, level=level, bpm=150.0,
                        offset_sec=0.12, density_bucket=dens, seed=7,
                        title="vendor-acceptance", wave="acceptance.ogg")
        charts = parse_tja(r["tja"])
        assert len(charts) == 1, f"{course}: expected one chart section"
        # chart_to_trace is the real acceptance: it rebuilds bar edges, tempo
        # and meter from the exported text and raises on anything inconsistent.
        trace = chart_to_trace(charts[0])
        assert len(trace.bars) == r["grid"]["n_bars"], (
            len(trace.bars), r["grid"]["n_bars"])
        n_notes = sum(c in "1234" for b in charts[0].bars for c in b.digits)
        n_spans = sum(c in "567" for b in charts[0].bars for c in b.digits)
        assert n_spans == r["n_spans"], (n_spans, r["n_spans"])

        assert r["grid"]["denom"] == 16, r["grid"]
        assert r["grid"]["triplets_representable"] is False
        assert r["grid"]["bar_denoms"] == "supplied", r["grid"]
        assert r["serving"]["min_onset_gap_sec"] == 0.02, r["serving"]
        assert r["serving"]["temperature"] == 1.0 and r["serving"]["top_p"] == 0.95
        assert r["serving_overrides"] == {}
        assert r["motif"]["verified"] is True, r["motif"]
        assert r["experimental"] is True
        assert n_notes == r["n_notes"], (n_notes, r["n_notes"])

        header = [ln for ln in r["tja"].splitlines() if ln.startswith("COURSE:")]
        assert header == [f"COURSE:{loader.TJA_COURSE_HEADER[course]}"], header
        assert charts[0].course == course, (charts[0].course, course)
        out[course] = {"n_notes": r["n_notes"], "n_spans": r["n_spans"],
                       "n_bars": r["grid"]["n_bars"], "header": header[0],
                       "reparsed_notes": n_notes}

    # same seed, same conditions, same chart
    a = sc.generate(_fake_mel(), "oni", level=8, bpm=150.0, offset_sec=0.12,
                    density_bucket=6, seed=7)["tja"]
    b = sc.generate(_fake_mel(), "oni", level=8, bpm=150.0, offset_sec=0.12,
                    density_bucket=6, seed=7)["tja"]
    assert a == b, "BarScript decode is not reproducible at a fixed seed"
    out["reproducible"] = True
    return out


# ---------------------------------------------------------------------------
# (c) the 1.x path is untouched


def test_c_legacy_library_untouched():
    lib = os.path.join(SPACES_DIR, "softchart")
    drift = {name: _md5(os.path.join(lib, name))
             for name in sorted(LEGACY_MD5)
             if _md5(os.path.join(lib, name)) != LEGACY_MD5[name]}
    assert not drift, (
        "1.x library modules changed since 2.0 was vendored in; the seven "
        f"published models are no longer a fixed baseline: {drift}")

    # Space-owned modules that may sit beside the 1.x library. The list is
    # explicit so a module arriving unnoticed is a failure, not a shrug:
    #   sc2_loader.py       the bridge to softchart.sc2
    #   barscript_grid.py   builds a BarScript bar timeline from a beat fit,
    #                       so a tempo change survives into the served grid
    # Neither is imported by any 1.x code path.
    added = sorted(n for n in os.listdir(lib)
                   if n.endswith(".py") and n not in LEGACY_MD5)
    assert added == ["barscript_grid.py", "sc2_loader.py"], added
    return {"unchanged": len(LEGACY_MD5), "added": added}


def test_c_legacy_api_unchanged():
    from softchart import generate as v1_generate
    from softchart import tja as v1_tja
    from softchart import vocab as v1_vocab

    mods = {"generate": v1_generate, "tja": v1_tja}
    bad = {}
    for key, want in LEGACY_SIGNATURES.items():
        mod, fn = key.split(".")
        got = str(inspect.signature(getattr(mods[mod], fn)))
        got = got.replace('"', "'")
        if got != want:
            bad[key] = got
    assert not bad, f"1.x public signatures changed: {bad}"

    got = {k: getattr(v1_vocab, k) for k in LEGACY_CONSTANTS if k != "VOCAB_SIZE"}
    got["VOCAB_SIZE"] = v1_vocab.VOCAB.size
    assert got == LEGACY_CONSTANTS, f"1.x constants changed: {got}"
    return got


def test_c_namespaces_isolated():
    """The two token tables coexist and neither can reach the other."""
    from softchart import vocab as v1_vocab
    from softchart.sc2 import vocab as v2_vocab

    assert v1_vocab is not v2_vocab
    assert v1_vocab.VOCAB.size == 1810 and v2_vocab.VOCAB.size == 2071
    assert v1_vocab.MAX_TGT == 1536 and v2_vocab.MAX_TGT == 960

    # every sc2 module resolves its imports inside the sc2 package
    from softchart.sc2 import generate as v2_generate
    from softchart.sc2 import model as v2_model

    assert os.path.abspath(v2_generate.__file__).startswith(SC2_DIR + os.sep)
    assert v2_generate.VOCAB is v2_vocab.VOCAB
    assert v2_model.__name__ == "softchart.sc2.model"

    from softchart import generate as v1_generate

    assert v1_generate.VOCAB is v1_vocab.VOCAB
    assert v1_generate is not v2_generate
    return {"v1_vocab": v1_vocab.VOCAB.size, "v2_vocab": v2_vocab.VOCAB.size}


def _legacy_decode():
    """One 1.x decode, in both the time and the slot regime app.py uses."""
    import numpy as np

    from softchart.generate import (generate_song, generate_song_slot, load_hf)

    model = load_hf(LEGACY_REPO, device="cpu")
    mel = _fake_mel()
    dur = mel.shape[1] / 86.1328125
    time_gen = generate_song(model, mel, "oni", level=8, density_bucket=7,
                             device="cpu", seed=0)
    bar = 240.0 / 150.0
    edges = np.arange(0.0, dur + bar, bar)
    grid = {"bpm": 150.0, "downbeats": edges[:-1], "bar": bar,
            "measure_edges": edges}
    slot_gen = generate_song_slot(model, mel, grid, "oni", level=8,
                                  density_bucket=7, device="cpu", seed=0)
    payload = {
        "time": [time_gen["hits"], time_gen["spans"]],
        "slot": [[list(map(_jsonable, h)) for h in slot_gen["hits_slots"]],
                 [list(map(_jsonable, s)) for s in slot_gen["spans_slots"]]],
    }
    blob = json.dumps(payload, sort_keys=True, default=_jsonable)
    return {"sha256": hashlib.sha256(blob.encode()).hexdigest(),
            "n_time_hits": len(time_gen["hits"]),
            "n_slot_hits": len(slot_gen["hits_slots"])}


def _jsonable(x):
    import numpy as np

    if isinstance(x, (np.integer,)):
        return int(x)
    if isinstance(x, (np.floating,)):
        return float(x)
    return x


def case_legacy_baseline():
    return _legacy_decode()


def case_legacy_after_sc2():
    """The same 1.x decode, after 2.0 is fully imported AND loaded.

    This is the regression the user's comparison depends on: 2.0 living in the
    same process must not move a single note of a 1.x chart.
    """
    from softchart import sc2_loader as loader

    sc = loader.load(device="cpu")
    sc.generate(_fake_mel(600), "oni", level=8, bpm=150.0, offset_sec=0.0,
                density_bucket=6, seed=1)
    out = _legacy_decode()
    out["sc2_release"] = sc.release_name
    return out


def _legacy_weights_cached():
    try:
        from huggingface_hub import try_to_load_from_cache
    except Exception:
        return False
    for name in ("model.safetensors", "config.json"):
        hit = try_to_load_from_cache(LEGACY_REPO, name)
        if not isinstance(hit, str):
            return False
    return True


def test_c_legacy_decode_unchanged():
    if not _legacy_weights_cached():
        return {"skipped": f"{LEGACY_REPO} is not in the local Hub cache"}
    before = _run_case("legacy-baseline")
    after = _run_case("legacy-after-sc2")
    assert before["sha256"] == after["sha256"], (
        "a 1.x decode changed once softchart.sc2 was imported: "
        f"{before} != {after}")
    return {"sha256": before["sha256"], "n_time_hits": before["n_time_hits"],
            "n_slot_hits": before["n_slot_hits"],
            "sc2_release": after["sc2_release"]}


# ---------------------------------------------------------------------------

CASES = {
    "package-only-load": case_package_only_load,
    "legacy-baseline": case_legacy_baseline,
    "legacy-after-sc2": case_legacy_after_sc2,
}

CHECKS = (
    ("(a) package-only load", test_a_package_only_load),
    ("(b) generate + re-parse", test_b_generate_and_reparse),
    ("(c) 1.x library untouched", test_c_legacy_library_untouched),
    ("(c) 1.x API unchanged", test_c_legacy_api_unchanged),
    ("(c) namespaces isolated", test_c_namespaces_isolated),
    ("(c) 1.x decode unchanged", test_c_legacy_decode_unchanged),
)


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--case", choices=sorted(CASES),
                    help="internal: run one subprocess case and print RESULT")
    args = ap.parse_args(argv)
    warnings.filterwarnings("ignore")
    os.environ.setdefault("HF_HUB_OFFLINE", "1")

    if args.case:
        print("RESULT " + json.dumps(CASES[args.case](), default=str))
        return 0

    failures = []
    for label, check in CHECKS:
        try:
            detail = check()
        except Exception as exc:  # noqa: BLE001 - a report, not a traceback
            failures.append(label)
            print(f"FAIL  {label}\n      {type(exc).__name__}: {exc}")
        else:
            mark = "SKIP" if isinstance(detail, dict) and "skipped" in detail \
                else "PASS"
            print(f"{mark}  {label}\n      "
                  + json.dumps(detail, default=str)[:600])
    print(("FAILED: " + ", ".join(failures)) if failures
          else "all checks passed")
    return 1 if failures else 0


if __name__ == "__main__":
    raise SystemExit(main())