code/make_multicommodity_manifests.py

"""
make_multicommodity_manifests.py — multi-commodity, multi-source benchmark
manifests for ARC-V Paper 2. Generalizes make_loso_manifests.py /
make_source_specific_manifests.py from one-commodity/two-sources to
C >= 2 commodities x S >= 3 acquisition sources.

Spec: paper/BENCHMARK_PROTOCOL.md. This script is the reference implementation.

It consumes a REGISTRY of the form

    {commodity: {"sources": {source: path_glob, ...},
                 "synonyms": {raw_label: canonical_name, ...}}}   # synonyms optional

and, for each commodity independently, emits in the existing manifest schema
(v2, additive over unified_benchmark.json):

  * manifest_<k>_src_<s>.json        — per-source (renumbered classes)
  * manifest_<k>_overlap.json        — classes present in >=2 sources
  * manifest_<k>_overlap_<s>.json    — overlap classes, one source (asym. probe)
  * manifest_<k>_loso_<held>.json    — train on others, test on held-out source

Source identity is carried EXPLICITLY in a per-sample `meta` block (1:1 with
`samples`), so v2-aware loaders never parse paths to recover the source. The
`samples` lists keep the Paper-1 [path, label] shape, so existing consumers
(src.dataset.load_manifest_splits / get_aifnet_dataloaders) still work.

Real images are NOT required to develop/test this: `--smoke` builds a synthetic
registry of dummy file lists and runs the whole pipeline on CPU.

    # smoke (no data needed, pure CPU):
    python make_multicommodity_manifests.py --smoke

    # real run from a registry file:
    python make_multicommodity_manifests.py --registry registry.json --val-frac 0.15

NOTE: this scaffold performs manifest construction only. The mandatory
pixel-verified de-duplication (BENCHMARK_PROTOCOL.md sec.5) reuses audit_dedup.py
and make_dedup_manifests.py unchanged and is invoked as a separate step; here we
only emit the `dedup` placeholder block and wire the hook (see run_dedup()).
"""
import argparse
import hashlib
import json
import re
from collections import defaultdict
from pathlib import Path

import numpy as np

# ── reuse existing helpers where they exist; fall back gracefully ───────────────
try:  # src.dataset.source_of is the Paper-1 v1 path->source fallback only
    from src.dataset import source_of as _v1_source_of  # noqa: F401
except Exception:  # pragma: no cover - keeps the smoke path importable anywhere
    _v1_source_of = None


# ── canonicalization ───────────────────────────────────────────────────────────
def slugify(label: str) -> str:
    """Canonical class name: lower, spaces/hyphens/dots -> underscore, collapse."""
    s = str(label).strip().lower()
    s = re.sub(r"[\s\-.]+", "_", s)
    s = re.sub(r"[^a-z0-9_]+", "", s)
    s = re.sub(r"_+", "_", s).strip("_")
    return s or "unknown"


def canonical(label: str, synonyms: dict) -> str:
    """Map a raw folder label to a canonical class name via synonyms then slugify."""
    if synonyms and label in synonyms:
        return synonyms[label]
    # also allow a synonym keyed by the slug
    sl = slugify(label)
    if synonyms and sl in synonyms:
        return synonyms[sl]
    return sl


# ── registry resolution: {commodity: [(path, source, canonical_class, raw), ...]} ─
def _class_label_from_path(path: str) -> str:
    """Raw class label = immediate parent directory name of the file."""
    return Path(str(path).replace("\\", "/")).parent.name


def resolve_registry(registry: dict, file_lists: dict | None = None) -> dict:
    """Expand a registry into per-commodity sample records.

    registry: {commodity: {"sources": {source: glob}, "synonyms": {...}}}
    file_lists: optional {commodity: {source: [paths...]}} that BYPASSES globbing
                (used by --smoke and by tests). When given, the registry glob is
                ignored for that (commodity, source) and these paths are used.

    Returns {commodity: list[(path, source, canonical_class, raw_label)]}.
    Raises if a real (non-overridden) source resolves to zero files.
    """
    out: dict[str, list] = {}
    for commodity, spec in registry.items():
        sources = spec.get("sources", {})
        synonyms = spec.get("synonyms", {})
        recs: list = []
        for source, glob in sources.items():
            paths = None
            if file_lists and commodity in file_lists and source in file_lists[commodity]:
                paths = list(file_lists[commodity][source])
            else:
                # real path: glob is relative to CWD (matches existing scripts)
                paths = [str(p) for p in sorted(Path().glob(glob))]
            if not paths:
                raise SystemExit(
                    f"[registry] commodity={commodity!r} source={source!r} glob={glob!r} "
                    f"matched 0 files; a silently-empty source corrupts LOSO."
                )
            for p in paths:
                raw = _class_label_from_path(p)
                recs.append((p, source, canonical(raw, synonyms), raw))
        out[commodity] = recs
    return out


def _registry_hash(registry: dict) -> str:
    blob = json.dumps(registry, sort_keys=True, default=str).encode()
    return "sha1:" + hashlib.sha1(blob).hexdigest()


# ── manifest assembly helpers ──────────────────────────────────────────────────
def _split_records(records, classes, val_frac, seed):
    """Carve train/val from a list of (path, source, class) over the given classes.

    Returns (train, val) where each is a list of (path, source, class). A
    per-(source,class) seeded shuffle puts `val_frac` into val. (Used for the
    training sources of a LOSO/per-source manifest; held-out test is added by caller.)
    """
    rng = np.random.default_rng(seed)
    by_sc = defaultdict(list)
    for p, s, c in records:
        if c in classes:
            by_sc[(s, c)].append(p)
    train, val = [], []
    for (s, c), paths in sorted(by_sc.items()):
        paths = list(paths)
        rng.shuffle(paths)
        n_val = max(1, int(round(len(paths) * val_frac))) if paths else 0
        for p in paths[:n_val]:
            val.append((p, s, c))
        for p in paths[n_val:]:
            train.append((p, s, c))
    return train, val


def _build_manifest(commodity, protocol, cls2idx, split_recs, *, seed, val_frac,
                    registry_hash, extra=None):
    """Assemble a v2 manifest dict from per-split (path, source, class) records.

    split_recs: {"train":[...], "val":[...], "test":[...]} of (path, source, class).
    cls2idx: canonical class name -> contiguous index for THIS manifest.
    """
    samples = {sp: [] for sp in ("train", "val", "test")}
    meta = {sp: [] for sp in ("train", "val", "test")}
    # per-class, per-split counts and the set of (source:raw-ish) provenance
    counts = {c: {"train": 0, "val": 0, "test": 0} for c in cls2idx}
    src_provenance = defaultdict(set)  # class -> {"source:class"}
    sources_seen = set()
    for sp in ("train", "val", "test"):
        for p, s, c in split_recs[sp]:
            idx = cls2idx[c]
            samples[sp].append([p, idx])
            meta[sp].append({"commodity": commodity, "source": s, "class": c})
            counts[c][sp] += 1
            src_provenance[c].add(f"{s}:{c}")
            sources_seen.add(s)

    classes = []
    for c in sorted(cls2idx, key=lambda x: cls2idx[x]):
        cnt = counts[c]
        classes.append({
            "name": c, "index": cls2idx[c], "commodity": commodity,
            "sources": sorted(src_provenance[c]),
            "total": cnt["train"] + cnt["val"] + cnt["test"],
            "train": cnt["train"], "val": cnt["val"], "test": cnt["test"],
        })

    source_map = {s: i for i, s in enumerate(sorted(sources_seen))}
    manifest = {
        "version": 2, "seed": seed, "protocol": protocol, "commodity": commodity,
        "splits": {"train": 1 - val_frac, "val": val_frac, "test": "held-out source"
                   if protocol == "leave-one-source-out" else val_frac},
        "classes": classes,
        "samples": samples,
        "meta": meta,
        "source_map": source_map,
        "registry_hash": registry_hash,
        # Placeholder; populated by the dedup step (audit_dedup / make_dedup_manifests).
        "dedup": {"method": "PENDING — run audit_dedup.py + make_dedup_manifests.py",
                  "cross_source_verified_dups": None,
                  "removed_from_test": 0,
                  "test_before": len(samples["test"]), "test_after": len(samples["test"])},
    }
    if extra:
        manifest.update(extra)
    return manifest


# ── the three manifest families, per commodity ─────────────────────────────────
def emit_for_commodity(commodity, records, *, outdir, val_frac, seed, registry_hash):
    """Emit per-source, overlap, and LOSO manifests for one commodity.

    records: list of (path, source, canonical_class, raw_label).
    Returns list of (Path, manifest_dict) written.
    """
    recs3 = [(p, s, c) for (p, s, c, _raw) in records]  # drop raw for assembly
    by_src_cls = defaultdict(lambda: defaultdict(list))
    for p, s, c in recs3:
        by_src_cls[s][c].append(p)
    sources = sorted(by_src_cls)
    written = []

    def _dump(name, manifest):
        out = Path(outdir) / name
        out.parent.mkdir(parents=True, exist_ok=True)
        json.dump(manifest, open(out, "w"), indent=2)
        written.append((out, manifest))
        nc = len(manifest["classes"])
        s = manifest["samples"]
        print(f"  {out.name}: {nc} classes | "
              f"train {len(s['train'])} val {len(s['val'])} test {len(s['test'])}")

    # classes present per source, and overlap (>=2 sources)
    classes_by_src = {s: set(by_src_cls[s]) for s in sources}
    all_classes = set().union(*classes_by_src.values()) if classes_by_src else set()
    overlap = sorted(c for c in all_classes
                     if sum(c in classes_by_src[s] for s in sources) >= 2)

    # (1) per-source manifests — renumber classes contiguously within the source
    for s in sources:
        scls = sorted(classes_by_src[s])
        cls2idx = {c: i for i, c in enumerate(scls)}
        tr, va = _split_records([(p, s, c) for (p, ss, c) in recs3 if ss == s],
                                set(scls), val_frac, seed)
        mf = _build_manifest(commodity, "source-specific", cls2idx,
                             {"train": tr, "val": va, "test": []},
                             seed=seed, val_frac=val_frac, registry_hash=registry_hash,
                             extra={"single_source": s})
        _dump(f"manifest_{commodity}_src_{s}.json", mf)

    # (2) overlap-only manifests (pooled + per-source asymmetric probes)
    if overlap:
        cls2idx = {c: i for i, c in enumerate(overlap)}
        tr, va = _split_records([r for r in recs3 if r[2] in overlap],
                                set(overlap), val_frac, seed)
        mf = _build_manifest(commodity, "overlap-only", cls2idx,
                             {"train": tr, "val": va, "test": []},
                             seed=seed, val_frac=val_frac, registry_hash=registry_hash)
        _dump(f"manifest_{commodity}_overlap.json", mf)
        for s in sources:
            srecs = [r for r in recs3 if r[1] == s and r[2] in overlap]
            if not srecs:
                continue
            tr, va = _split_records(srecs, set(overlap), val_frac, seed)
            mf = _build_manifest(commodity, "overlap-only", cls2idx,
                                 {"train": tr, "val": va, "test": []},
                                 seed=seed, val_frac=val_frac,
                                 registry_hash=registry_hash, extra={"single_source": s})
            _dump(f"manifest_{commodity}_overlap_{s}.json", mf)
    else:
        print(f"  [overlap] commodity={commodity!r}: no class shared by >=2 sources")

    # (3) Leave-One-Source-Out — train on others, test on the entire held-out source
    if len(sources) < 2:
        print(f"  [loso] commodity={commodity!r}: need >=2 sources, have {len(sources)}")
        return written
    for held in sources:
        train_srcs = [s for s in sources if s != held]
        train_classes = set().union(*[classes_by_src[s] for s in train_srcs])
        common = sorted(set(classes_by_src[held]) & train_classes)
        if not common:
            print(f"  [loso/skip] held={held!r}: no class shared with training sources")
            continue
        cls2idx = {c: i for i, c in enumerate(common)}
        tr, va = _split_records([(p, s, c) for (p, s, c) in recs3
                                 if s in train_srcs and c in common],
                                set(common), val_frac, seed)
        test = [(p, held, c) for c in common for p in by_src_cls[held][c]]
        mf = _build_manifest(commodity, "leave-one-source-out", cls2idx,
                             {"train": tr, "val": va, "test": test},
                             seed=seed, val_frac=val_frac, registry_hash=registry_hash,
                             extra={"held_out_source": held, "train_sources": train_srcs,
                                    "degenerate_loso": len(sources) < 3})
        _dump(f"manifest_{commodity}_loso_{held}.json", mf)

    if len(sources) < 3:
        print(f"  [warn] commodity={commodity!r}: only {len(sources)} sources; LOSO "
              f"manifests flagged degenerate_loso=true (protocol requires >=3).")
    return written


# ── dedup hook (reuses the existing, validated tooling — NOT reimplemented) ──────
def run_dedup(written):
    """Wire the mandatory pixel-verified dedup step (BENCHMARK_PROTOCOL.md sec.5).

    This intentionally does not duplicate audit_dedup.py. For each emitted manifest
    it shells out conceptually to make_dedup_manifests.py (train/test leakage) and,
    per commodity, to audit_dedup.py (cross-source hard gate). Kept as a documented
    stub so the smoke path stays pure/dependency-light (no PIL, no real images).
    """
    print("\n[dedup] mandatory pixel-verified de-duplication is a separate step:")
    print("  per manifest : python make_dedup_manifests.py --manifest <file>   # (L) train/test leakage")
    print("  per commodity: python audit_dedup.py --base <overlap_manifest>    # (X) cross-source HARD gate (must be 0)")
    print(f"  ({len(written)} manifests emitted; dedup blocks are PENDING until run.)")


# ── synthetic registry for the CPU smoke path (no images needed) ────────────────
def synthetic_registry_and_files(seed=0):
    """Build a fake multi-commodity, multi-source registry + dummy file lists.

    2 commodities x 3 sources, with deliberate overlap and source-specific classes,
    so every manifest family (per-source, overlap, LOSO) is exercised. Paths are
    synthetic strings; no file ever needs to exist.
    """
    rng = np.random.default_rng(seed)

    def files(commodity, source, classes, n_per):
        fl = []
        for c in classes:
            for i in range(n_per + int(rng.integers(0, 3))):
                fl.append(f"_smoke/{commodity}/{source}/{c}/img_{i:03d}.jpg")
        return fl

    # rice: 3 sources, shared {basmati, jasmine}, plus per-source extras + a synonym
    rice_files = {
        "koklu":       files("rice", "koklu",       ["Basmati", "Jasmine", "Arborio"], 12),
        "bangladeshi": files("rice", "bangladeshi", ["basmati_long", "Jasmine", "Miniket"], 10),
        "bina":        files("rice", "bina",        ["Basmati", "Jasmine", "BR29"], 8),
    }
    # spices: 3 sources, shared {black_pepper, cumin}, plus extras
    spice_files = {
        "indian":         files("spices", "indian",         ["Black Pepper", "Cumin", "Clove"], 11),
        "spice_spectrum": files("spices", "spice_spectrum", ["black pepper", "cumin", "Cardamom"], 9),
        "market":         files("spices", "market",         ["Black Pepper", "Cumin", "Fenugreek"], 7),
    }
    registry = {
        "rice": {
            "sources": {k: f"_smoke/rice/{k}/*/*.jpg" for k in rice_files},
            "synonyms": {"basmati_long": "basmati"},   # canonicalize to 'basmati'
        },
        "spices": {
            "sources": {k: f"_smoke/spices/{k}/*/*.jpg" for k in spice_files},
            "synonyms": {},                            # slugify handles 'Black Pepper'
        },
    }
    file_lists = {"rice": rice_files, "spices": spice_files}
    return registry, file_lists


# ── self-check for the smoke path ───────────────────────────────────────────────
def _validate(manifest):
    """Assert the v2 invariants from BENCHMARK_PROTOCOL.md sec.2."""
    for sp in ("train", "val", "test"):
        assert len(manifest["meta"][sp]) == len(manifest["samples"][sp]), \
            f"meta/samples length mismatch in {sp}"
    idx2name = {c["index"]: c["name"] for c in manifest["classes"]}
    # contiguous 0..k-1
    idxs = sorted(c["index"] for c in manifest["classes"])
    assert idxs == list(range(len(idxs))), f"non-contiguous class indices: {idxs}"
    for sp in ("train", "val", "test"):
        for (p, lbl), mt in zip(manifest["samples"][sp], manifest["meta"][sp]):
            assert idx2name[lbl] == mt["class"], "label<->class mismatch"
            assert mt["commodity"] == manifest["commodity"] or manifest["commodity"] == "*"


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--registry", help="path to registry JSON "
                    "{commodity:{sources:{src:glob},synonyms:{...}}}")
    ap.add_argument("--smoke", action="store_true",
                    help="run the CPU smoke path on a synthetic registry (no images)")
    ap.add_argument("--val-frac", type=float, default=0.15)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--outdir", default="outputs")
    ap.add_argument("--no-dedup-hint", action="store_true",
                    help="suppress the dedup next-step reminder")
    args = ap.parse_args()

    file_lists = None
    if args.smoke:
        registry, file_lists = synthetic_registry_and_files(seed=args.seed)
        if not args.outdir or args.outdir == "outputs":
            args.outdir = "outputs/_smoke_manifests"
        print(f"[smoke] synthetic registry: commodities={list(registry)}; "
              f"outdir={args.outdir}")
    elif args.registry:
        registry = json.load(open(args.registry))
    else:
        ap.error("provide --registry <file> or --smoke")

    registry_hash = _registry_hash(registry)
    resolved = resolve_registry(registry, file_lists=file_lists)

    all_written = []
    for commodity, records in resolved.items():
        srcs = sorted({s for _, s, _, _ in records})
        ncls = len({c for _, _, c, _ in records})
        print(f"\ncommodity={commodity!r}: {len(records)} samples, "
              f"sources={srcs}, classes={ncls}")
        written = emit_for_commodity(
            commodity, records, outdir=args.outdir,
            val_frac=args.val_frac, seed=args.seed, registry_hash=registry_hash)
        all_written += written

    if args.smoke:
        for _path, mf in all_written:
            _validate(mf)
        print(f"\n[smoke] OK — {len(all_written)} manifests written & validated "
              f"(v2 invariants hold).")

    if not args.no_dedup_hint:
        run_dedup(all_written)

    print(f"\nwrote {len(all_written)} manifests to {args.outdir}.")


if __name__ == "__main__":
    main()