#!/usr/bin/env python3
"""Fetch and lay out the benchmark sample corpus.

audiofiles-bench times the analysis pipeline over a corpus laid out as
samples/training/<class>/ and samples/test-suite/<kind>/. The per-class folders
are also the filename-derived ground truth the browse-axes measurements used.
That tree is gitignored and has to be rebuilt per machine. This script fetches
the source datasets and maps them into that layout.

The large datasets are opt-in because they are tens of gigabytes and, for
NSynth, 16 kHz mono, which skews per-file decode timings away from what a real
44.1/48 kHz library costs.

Nothing here is licence-clean as a whole. Two of the five mix per-sound
licences, including CC-BY-NC, so the registry's `license` field describes the
mix rather than claiming one label. See summarize_fsl10k_licenses.

    ./scripts/corpus.py --list
    ./scripts/corpus.py --dest /media/max/T9/af-corpus
    ./scripts/corpus.py --dest /media/max/T9/af-corpus --datasets nsynth

Datasets land in <dest>/_downloads (archives), <dest>/_raw (extracted), and
<dest>/samples (the layout the bench reads). Point the bench at it with
AF_BENCH_CORPUS=<dest>/samples.

Every source dataset is wav-only, so the per-format decode section of the bench
gets its non-wav files from a transcode step (needs ffmpeg on PATH) rather than
from a download. See build_formats.
"""

import argparse
import json
import shutil
import subprocess
import sys
import urllib.request
from pathlib import Path

# Dataset registry.
#
# `default` marks the ones that run without --datasets. It is a size-and-time
# judgement and says nothing about licence: fsl10k is a default and a sixth of
# it is CC-BY-NC. The two questions were conflated here until 2026-08-08, when
# the entry still claimed a flat "CC-BY 4.0" for a corpus with four licences in
# it.
#
# Nothing derived from any of this ships today. The bundled `.afcl` layer that
# would have was retired on 2026-08-08 (docs/ml_classifier.md, "What ships in
# the binary, and under what licence"), so every dataset here is measurement
# material only, which the NC clips are fine for. The licence is recorded
# accurately anyway: that position has already flipped once, and reconstructing
# provenance after the fact is what this file exists to avoid.
DATASETS = {
    "reverb-drums": {
        "default": True,
        "size": "570 MB",
        "count": "1,786 files",
        "license": "CC-BY 4.0",
        "url": "https://archive.org/download/reverb-drum-machines-complete-collection/"
        "Reverb%20Drum%20Machines%20_%20The%20Complete%20Collection.7z",
        "archive": "reverb-drums.7z",
        "desc": "Reverb drum machine packs. Instrument in the filename; the "
        "only default source of labeled kick/snare/hihat one-shots.",
    },
    "fsl10k": {
        "default": True,
        "size": "8.8 GB",
        "count": "9,455 loops",
        "license": "per sound: CC0, CC-BY, CC-BY-NC, Sampling+",
        "url": "https://zenodo.org/api/records/3967852/files/FSL10K.zip/content",
        "archive": "FSL10K.zip",
        "extra": {
            "annotations.zip": "https://zenodo.org/api/records/3967852/files/annotations.zip/content"
        },
        "desc": "Freesound Loop Dataset. The only source here with ground-truth "
        "tempo and key, so it is what BPM/key accuracy can be scored against. "
        "Zenodo publishes it as CC-BY; the sounds are not. 1,436 of 9,493 are "
        "NC or Sampling+: fine for measurement, not for anything that ships. "
        "Select on metadata.json, never on this entry.",
    },
    "nsynth": {
        "default": False,
        "size": "~30 GB",
        "count": "305,979 notes",
        "license": "CC-BY 4.0",
        "url": "http://download.magenta.tensorflow.org/datasets/nsynth/nsynth-train.jsonwav.tar.gz",
        "archive": "nsynth-train.jsonwav.tar.gz",
        "desc": "NSynth. Use for count-scale (300k rows, 300k blobs in one flat "
        "dir). 16 kHz mono, so do not mix its timings into a throughput number.",
    },
    "fsd50k": {
        "default": False,
        "size": "~30 GB",
        "count": "51,197 clips",
        "license": "CC-BY (per-clip varies, includes CC-BY-NC)",
        "url": "https://zenodo.org/api/records/4060432/files/FSD50K.dev_audio.zip/content",
        "archive": "FSD50K.dev_audio.zip",
        "desc": "FSD50K, 200 AudioSet classes. Contains CC-BY-NC clips: fine for "
        "measurement, not for anything that ships.",
    },
    "percussive": {
        "default": False,
        "size": "119 MB",
        "count": "10,254 sounds",
        "license": "Attribution",
        "url": "https://zenodo.org/api/records/3665275/files/one_shot_percussive_sounds.zip/content",
        "archive": "one_shot_percussive_sounds.zip",
        "desc": "Freesound one-shot percussion. No class labels and 16 kHz "
        "normalized to 1 s, so it is bulk filler, not accuracy material.",
    },
}

# Freesound licence URLs as they appear in FSL10K's metadata.json, mapped to a
# short name and whether a derived work under that licence could ship.
#
# Sampling+ is counted as cannot-ship rather than argued about: it permits
# sampling into a new work but bars advertising use and verbatim
# redistribution, which is ambiguous for a corpus-derived artifact and worth
# 221 sounds out of 9,493.
FSL10K_LICENSES = {
    "http://creativecommons.org/publicdomain/zero/1.0/": ("CC0 1.0", True),
    "http://creativecommons.org/licenses/by/3.0/": ("CC-BY 3.0", True),
    "http://creativecommons.org/licenses/by-nc/3.0/": ("CC-BY-NC 3.0", False),
    "http://creativecommons.org/licenses/sampling+/1.0/": ("Sampling+ 1.0", False),
}


def summarize_fsl10k_licenses(raw: Path) -> dict | None:
    """Count FSL10K's per-sound licences off its metadata.json.

    The dataset-level label cannot answer the licence question for this corpus
    and the registry no longer pretends it can, so the real mix is measured on
    every run and recorded in the manifest. Anything that later selects sounds
    out of FSL10K reads this, or reads metadata.json itself; the point is that
    it does not read a single label off the registry entry.

    An unrecognised licence URL counts as cannot-ship and is named in the
    result. Defaulting the other way would let a new Freesound licence into a
    shippable set silently, which is the exact failure this function exists to
    close.

    Returns None when the metadata is not there -- --no-build runs and partial
    extractions are normal, and a missing count is better than a wrong one.
    """
    meta = raw / "metadata.json"
    if not meta.is_file():
        return None
    try:
        sounds = json.loads(meta.read_text())
    except (OSError, json.JSONDecodeError) as e:
        print(f"  warning: cannot read {meta} ({e}); licence mix not recorded")
        return None
    if not isinstance(sounds, dict):
        return None

    counts: dict[str, int] = {}
    unrecognised: dict[str, int] = {}
    shippable = 0
    for sound in sounds.values():
        url = sound.get("license") if isinstance(sound, dict) else None
        name, ok = FSL10K_LICENSES.get(url, (None, False))
        if name is None:
            name = str(url)
            unrecognised[name] = unrecognised.get(name, 0) + 1
        counts[name] = counts.get(name, 0) + 1
        shippable += ok

    summary = {
        "total": len(sounds),
        "by_license": dict(sorted(counts.items())),
        "shippable": shippable,
        "excluded": len(sounds) - shippable,
        "source": "metadata.json, per sound",
    }
    if unrecognised:
        summary["unrecognised"] = dict(sorted(unrecognised.items()))
    return summary


# Filename/dirname keyword -> bench training class.
#
# Ordered, first match wins. Order is load-bearing: "loop" has to be tested
# before any instrument keyword because loop files in these packs are named
# "<machine> Loop3.wav" with no instrument in the name at all, and several
# instrument keywords are substrings of each other.
CLASS_RULES = [
    (["kick", "bassdrum", "bass drum", " bd", "_bd", "kik"], "kick"),
    (["snare", " sd", "_sd", "rimshot", "rim shot", " rim", "_rim",
      "sidestick", "side stick", "stick"], "snare"),
    (["hihat", "hi hat", "hi-hat", "chh", "ohh", " hh", "_hh", "hat"], "hihat"),
    (["cymbal", "crash", "ride", "splash", "china", "gong"], "cymbal"),
    (["clap", "handclap", " cp", "_cp"], "clap"),
    (["tom"], "tom"),
    (
        [
            "perc", "cowbell", "clave", "clava", "maraca", "tambor", "tambour",
            "bongo", "conga", "guiro", "shaker", "shake", "block", "triangle",
            "agogo", "cabasa", "timbale", "whistle", "bell", "scratch",
            "chime", "click", "beep", "quijada", "steel drum",
        ],
        "percussion",
    ),
]

TRAINING_CLASSES = ["kick", "snare", "hihat", "cymbal", "clap", "tom", "percussion"]

# The one dataset that fills samples/training/, and so the one the labels are
# derived from. Recorded in the manifest as `training_source` because the
# `datasets` list answers a different question -- what was fetched -- and the
# two diverge as soon as anything else is downloaded into the same corpus.
TRAINING_DATASET = "reverb-drums"


def classify_name(path: Path) -> tuple[str | None, str]:
    """Map a source file to a training class from its name and parent dirs.

    Returns (class, status) where status is one of "ok", "loop", "ambiguous",
    "unlabeled".

    Matches against the filename plus its two parent directories, since these
    packs sometimes put the instrument in a subdirectory ("..._Tom/") and
    sometimes only in the filename ("... Tom3.wav").

    Files matching more than one class are dropped rather than resolved by rule
    order. These packs contain genuinely ambiguous names -- "Kick_Cowbell.wav",
    "Tom-Cymbal.wav" -- and first-match-wins would silently assign one of the
    two at random. A wrong label is worse than a missing one here: these folders
    are the ground truth the browse-axes measurements are read against, so a
    mislabeled file looks like a real result forever.
    """
    hay = " ".join([path.name, path.parent.name, path.parent.parent.name]).lower()

    # Loops carry no instrument in the name at all in these packs, so they are
    # checked first and routed away from the training set entirely.
    if "loop" in hay:
        return None, "loop"

    matched = {cls for keywords, cls in CLASS_RULES if any(k in hay for k in keywords)}
    if len(matched) > 1:
        return None, "ambiguous"
    if len(matched) == 1:
        return matched.pop(), "ok"
    return None, "unlabeled"


def run(cmd: list[str]) -> None:
    subprocess.run(cmd, check=True)


def download(url: str, dest: Path) -> None:
    if dest.exists() and dest.stat().st_size > 0:
        print(f"  have {dest.name} ({dest.stat().st_size / 1e9:.2f} GB), skipping")
        return
    print(f"  downloading {dest.name} ...")
    tmp = dest.with_suffix(dest.suffix + ".part")
    urllib.request.urlretrieve(url, tmp)
    tmp.rename(dest)
    print(f"  got {dest.name} ({dest.stat().st_size / 1e9:.2f} GB)")


def extract(archive: Path, into: Path, marker: Path | None = None) -> None:
    """Extract `archive` into `into`.

    `marker` is the path whose existence means this archive is already
    unpacked. It defaults to "`into` is non-empty", which is right for the first
    archive into a directory but wrong for a second one (annotations landing
    next to audio), where the directory is already full.
    """
    done = marker.exists() if marker is not None else (into.exists() and any(into.iterdir()))
    if done:
        print(f"  already extracted {archive.name}, skipping")
        return
    into.mkdir(parents=True, exist_ok=True)
    print(f"  extracting {archive.name} ...")
    name = archive.name.lower()
    if name.endswith(".7z"):
        run(["7z", "x", "-y", f"-o{into}", str(archive)])
    elif name.endswith(".zip"):
        run(["unzip", "-q", "-o", str(archive), "-d", str(into)])
    elif name.endswith((".tar.gz", ".tgz")):
        run(["tar", "xzf", str(archive), "-C", str(into)])
    else:
        sys.exit(f"unknown archive type: {archive.name}")


def build_training(raw: Path, samples: Path) -> dict[str, int]:
    """Copy labeled one-shots into samples/training/<class>/.

    Copies rather than symlinks: the corpus is expected to live on the exFAT
    test drive, which has no symlink support at all.
    """
    counts: dict[str, int] = {c: 0 for c in TRAINING_CLASSES}
    counts["_loops"] = 0
    counts["_ambiguous"] = 0
    counts["_unlabeled"] = 0

    training = samples / "training"
    loops = samples / "test-suite" / "genres" / "loops"
    for c in TRAINING_CLASSES:
        (training / c).mkdir(parents=True, exist_ok=True)
    loops.mkdir(parents=True, exist_ok=True)

    dropped: list[str] = []

    for src in sorted(raw.rglob("*")):
        if not src.is_file() or src.suffix.lower() not in (".wav", ".aif", ".aiff", ".flac"):
            continue
        cls, status = classify_name(src)
        if status == "loop":
            dest_dir, key = loops, "_loops"
        elif status == "ok":
            dest_dir, key = training / cls, cls
        else:
            counts[f"_{status}"] += 1
            dropped.append(f"{status}: {src.name}")
            continue

        # Flatten with a pack-qualified name so same-named files across packs
        # do not collide (every pack has a "Kick.wav").
        flat = f"{src.parent.parent.name}__{src.name}".replace("/", "_")
        dest = dest_dir / flat
        if not dest.exists():
            shutil.copy2(src, dest)
        counts[key] += 1

    # Written out rather than just counted: when a measurement looks off, the
    # first question is always whether the corpus or the code is wrong.
    (samples / "DROPPED.txt").write_text("\n".join(sorted(dropped)) + "\n")

    return counts


# Encoder settings per format, as ffmpeg output arguments.
#
# Every arm is re-encoded, including wav. Copying the source wav instead would
# leave that one row measuring a different signal shape from the other three,
# since the source packs mix sample rates, bit depths, and channel counts. All
# four are normalized to the same 44.1 kHz stereo PCM so the only thing that
# differs between rows is the container and codec.
FORMAT_ENCODERS = {
    "wav": ["-c:a", "pcm_s16le"],
    "aiff": ["-c:a", "pcm_s16be"],
    "flac": ["-c:a", "flac"],
    # Lossy by nature, so its decoded output is not bit-identical to the other
    # three. That is inherent to comparing codecs and does not affect the
    # timing, which is what this corpus exists to measure.
    "mp3": ["-c:a", "libmp3lame", "-b:a", "320k"],
}


def build_formats(samples: Path, count: int) -> dict[str, int]:
    """Transcode a fixed subset of loops into samples/test-suite/formats/<ext>/.

    Section 2 of the bench times decode per format, which needs the same audio
    in every format. Nothing in the source datasets provides that: every corpus
    file, across all three datasets, is wav. Pointing the section at unrelated
    per-format files instead would measure content rather than codec.

    Sourced from the loops rather than the one-shots because the one-shots run
    about 0.2s, short enough that open-and-probe overhead dominates and the
    codec difference disappears into it. Loops run about 2.5s.
    """
    src_dir = samples / "test-suite" / "genres" / "loops"
    sources = sorted(p for p in src_dir.glob("*.wav") if p.is_file())
    if not sources:
        print(f"  no loops under {src_dir}, skipping")
        return {}

    # Stride rather than head: the flattened names sort by pack, so the first N
    # would all come from one drum machine and share its recording character.
    if len(sources) > count:
        stride = len(sources) / count
        sources = [sources[int(i * stride)] for i in range(count)]

    out_root = samples / "test-suite" / "formats"
    for ext in FORMAT_ENCODERS:
        (out_root / ext).mkdir(parents=True, exist_ok=True)

    counts = {ext: 0 for ext in FORMAT_ENCODERS}
    failed = 0
    for src in sources:
        outputs = {ext: out_root / ext / f"{src.stem}.{ext}" for ext in FORMAT_ENCODERS}
        if all(p.exists() for p in outputs.values()):
            for ext in FORMAT_ENCODERS:
                counts[ext] += 1
            continue

        written: list[Path] = []
        for ext, enc in FORMAT_ENCODERS.items():
            dest = outputs[ext]
            cmd = ["ffmpeg", "-v", "error", "-y", "-i", str(src),
                   "-ac", "2", "-ar", "44100", *enc, str(dest)]
            try:
                subprocess.run(cmd, check=True)
            except subprocess.CalledProcessError:
                break
            written.append(dest)
        else:
            for ext in FORMAT_ENCODERS:
                counts[ext] += 1
            continue

        # All-or-nothing. A file present in three formats but not the fourth
        # would silently change which audio each row averages over, so the
        # per-format means would no longer be comparable.
        for p in written:
            p.unlink(missing_ok=True)
        failed += 1

    if failed:
        print(f"  dropped {failed} source file(s) that would not encode to all formats")
    return counts


def load_manifest(path: Path) -> dict:
    """Read the existing manifest, or start an empty one.

    A run fetches the datasets it was asked for and rebuilds only what those
    datasets feed, so it knows about a slice of the corpus rather than all of
    it. Writing a manifest built from that slice alone is how the attribution
    got lost: a `--datasets nsynth` run left samples/training/ full of
    reverb-drums files and replaced the credit with a dataset the labels never
    came from. Both are CC-BY 4.0 so the licence class survived it, but CC-BY
    asks for credit to the work actually used. Merge instead.

    A corrupt manifest is not fatal: this run is about to rewrite it with
    whatever it knows, and refusing to proceed would leave the bad file in
    place. The old contents are lost, which is the reason for the warning.
    """
    if not path.exists():
        return {}
    try:
        doc = json.loads(path.read_text())
    except (OSError, json.JSONDecodeError) as e:
        print(f"  warning: cannot read {path} ({e}); starting a fresh manifest")
        return {}
    return doc if isinstance(doc, dict) else {}


def merge_datasets(old: list, new: list) -> list:
    """Union two dataset lists by name, with this run's entry winning.

    This run's entry wins because it was just fetched from the URL it names,
    whereas the recorded one may predate a registry edit. Sorted so a re-run
    over the same corpus produces a byte-identical manifest.
    """
    by_name = {d["name"]: d for d in old if isinstance(d, dict) and "name" in d}
    for d in new:
        by_name[d["name"]] = d
    return [by_name[k] for k in sorted(by_name)]


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--dest", type=Path, default=Path("/media/max/T9/af-corpus"))
    ap.add_argument("--datasets", help="comma-separated; default = the small ones")
    ap.add_argument("--list", action="store_true", help="show the registry and exit")
    ap.add_argument("--no-build", action="store_true", help="fetch and extract only")
    ap.add_argument("--formats-count", type=int, default=120,
                    help="loops to transcode into each format for the decode bench")
    args = ap.parse_args()

    if args.list:
        print(f"{'dataset':<14} {'size':>8}  {'default':<8} {'license':<46} count")
        for name, d in DATASETS.items():
            print(
                f"{name:<14} {d['size']:>8}  {str(d['default']):<8} "
                f"{d['license']:<46} {d['count']}"
            )
            print(f"{'':>14} {d['desc']}")
        return

    if args.datasets:
        wanted = [d.strip() for d in args.datasets.split(",")]
        unknown = [d for d in wanted if d not in DATASETS]
        if unknown:
            sys.exit(f"unknown dataset(s): {', '.join(unknown)}")
    else:
        wanted = [n for n, d in DATASETS.items() if d["default"]]

    downloads = args.dest / "_downloads"
    raw = args.dest / "_raw"
    samples = args.dest / "samples"
    downloads.mkdir(parents=True, exist_ok=True)

    manifest_path = args.dest / "MANIFEST.json"
    manifest = load_manifest(manifest_path)
    manifest["layout"] = str(samples)
    fetched: list[dict] = []

    for name in wanted:
        d = DATASETS[name]
        print(f"\n=== {name} ({d['size']}, {d['license']}) ===")
        archive = downloads / d["archive"]
        download(d["url"], archive)
        extract(archive, raw / name)
        # Extras unpack alongside the main archive, not into their own tree:
        # the accuracy bench expects annotations/ next to audio/ under one root.
        for extra_name, extra_url in d.get("extra", {}).items():
            extra_path = downloads / extra_name
            download(extra_url, extra_path)
            if extra_path.suffix.lower() in (".zip", ".7z", ".gz"):
                extract(extra_path, raw / name, marker=raw / name / Path(extra_name).stem)

        entry = {"name": name, "license": d["license"], "source": d["url"]}
        # FSL10K is the one dataset here shipping a per-sound licence field, so
        # it is the one where the manifest can carry something better than the
        # registry's prose. Measured rather than copied: the registry entry is
        # what was wrong in the first place.
        if name == "fsl10k":
            summary = summarize_fsl10k_licenses(raw / name)
            if summary is not None:
                entry["per_sound_licenses"] = summary
                print(
                    f"  licences: {summary['shippable']} shippable, "
                    f"{summary['excluded']} excluded, of {summary['total']}"
                )
                for lic, n in summary["by_license"].items():
                    print(f"    {lic:<16} {n}")
        fetched.append(entry)

    manifest["datasets"] = merge_datasets(manifest.get("datasets", []), fetched)

    if not args.no_build and TRAINING_DATASET in wanted:
        print("\n=== building training layout ===")
        counts = build_training(raw / TRAINING_DATASET, samples)
        for k, v in counts.items():
            print(f"  {k:<14} {v}")
        manifest["training_counts"] = counts
        # Names the dataset the labels came from, which is what anything
        # shipping those labels would have to credit. Nothing reads it as of
        # 2026-08-08: the bundled .afcl generator that required it was retired
        # with the layer itself. Kept because it costs a line and the credit is
        # unreconstructable once the corpus is rebuilt with other datasets in
        # it, which is how it got lost the first time (see load_manifest).
        manifest["training_source"] = TRAINING_DATASET

    # Runs off the built layout, not off a dataset, so it is gated on the loops
    # existing rather than on which datasets were requested.
    if not args.no_build and (samples / "test-suite" / "genres" / "loops").exists():
        print("\n=== building format corpus ===")
        if shutil.which("ffmpeg") is None:
            print("  ffmpeg not found, skipping (section 2 of the bench will report no files)")
        else:
            fmt_counts = build_formats(samples, args.formats_count)
            for k, v in fmt_counts.items():
                print(f"  {k:<14} {v}")
            # Recorded as derived so the license question does not come up
            # again: these are transcodes of corpus files, not a new dataset.
            manifest["format_counts"] = fmt_counts
            manifest["format_corpus"] = "transcoded from test-suite/genres/loops"

    # Provenance matters here: the corpus mixes licenses, and anything that
    # feeds a model needs that recorded rather than reconstructed later.
    manifest_path.write_text(json.dumps(manifest, indent=2))
    print(f"\nwrote {manifest_path}")
    print(f"point the bench at it:  export AF_BENCH_CORPUS={samples}")


if __name__ == "__main__":
    main()