tools/build_electronic_reference.py

"""build_electronic_reference.py — Seed the MERT reference DB with electronic
gear (drum machines + synths) so embedding-based ID works on electronic audio.

Backs up the existing (acoustic) DB first, then adds several variants per gear
class (different tempos/seeds for drums, different pitches for synths) for
nearest-neighbor robustness.

Usage: python3 tools/build_electronic_reference.py
"""
import logging
import os
import shutil
import sys
import tempfile

import numpy as np
import soundfile as sf

logging.basicConfig(level=logging.WARNING)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from backend.mert_embeddings import MERTExtractor, get_reference_db, InstrumentReferenceDB
from tools.synth_gear import SR, KITS, SYNTHS, synth_riff

DRUM_SEEDS = [1, 2, 3]                            # pattern/tempo variants per machine
# Dense root coverage (~G1..D4, every few semitones) so a query at any pitch has
# a near reference. Synths are rendered as arpeggio riffs, not sustained drones.
SYNTH_ROOTS = [49, 65, 82, 98, 123, 147, 185, 233, 294]


def _embed_array(extractor, y):
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tf:
        path = tf.name
    try:
        sf.write(path, y, SR)
        return extractor.extract(path)
    finally:
        os.unlink(path)


def main():
    db = get_reference_db()
    # Back up the acoustic-only DB once.
    backup = InstrumentReferenceDB.DB_PATH.with_suffix(".backup_acoustic.json")
    if not backup.exists():
        shutil.copy(InstrumentReferenceDB.DB_PATH, backup)
        print(f"Backed up acoustic DB -> {backup.name}")

    # Don't double-seed: skip if electronic entries already present.
    existing_models = {e["model"] for e in db.entries}
    already = existing_models & (set(KITS) | set(SYNTHS))
    if already:
        print(f"Electronic entries already present ({len(already)} models). "
              f"Delete them or restore the backup to rebuild. Aborting.")
        return

    start = len(db.entries)
    extractor = MERTExtractor()
    print(f"Seeding electronic gear (starting from {start} acoustic entries)…\n")

    # Drum machines
    for model, render in KITS.items():
        for seed in DRUM_SEEDS:
            y = render(seed=seed)
            emb = _embed_array(extractor, y)
            db.add(emb, model=model, family="drum_machine",
                   source_file=f"synth:{model}:seed{seed}")
        print(f"  ✓ {model}  ({len(DRUM_SEEDS)} variants)")

    # Synths — rendered as arpeggio riffs across dense root coverage
    for model, (family, fn) in SYNTHS.items():
        for root in SYNTH_ROOTS:
            y = synth_riff(fn, root=root)
            emb = _embed_array(extractor, y)
            db.add(emb, model=model, family=family,
                   source_file=f"synth_riff:{model}:{root}Hz")
        print(f"  ✓ {model}  ({len(SYNTH_ROOTS)} roots, family={family})")

    added = len(db.entries) - start
    print(f"\nDone. Added {added} electronic reference embeddings "
          f"(DB now {len(db.entries)} total).")


if __name__ == "__main__":
    main()