"""
Build source-specific manifests from the unified benchmark.

Outputs three new manifests:
  outputs/manifest_ss_only.json      — SpiceSpectrum samples only (11 classes)
  outputs/manifest_indian_only.json  — Mendeley Indian samples only (19 classes)
  outputs/manifest_overlap_only.json — only the 8 classes present in both sources

Each preserves the same 70/15/15 split structure as the unified manifest.
Class indices are RENUMBERED contiguously within each manifest (so SS-only has
classes 0..10, Indian-only 0..18, overlap 0..7).
"""
import json
from pathlib import Path

ROOT = Path(__file__).parent
UNIFIED = ROOT / "outputs" / "unified_benchmark.json"
OUT_DIR = ROOT / "outputs"


def _source_of(path: str) -> str:
    p = path.replace("\\", "/").lower()
    if "/spice_spectrum/" in p:
        return "spice_spectrum"
    if "/indian_spices/" in p:
        return "indian"
    return "unknown"


def _filter_manifest(uni, keep_source: str = None, keep_overlap_only: bool = False):
    """Return a new manifest containing only samples matching the filter."""
    out = {"version": uni["version"], "seed": uni["seed"],
           "splits": uni["splits"], "classes": [], "samples": {"train": [], "val": [], "test": []}}
    # Determine which classes to keep
    overlap_names = [c["name"] for c in uni["classes"] if len(c["sources"]) > 1]
    if keep_overlap_only and keep_source == "spice_spectrum":
        keep_class_names = overlap_names
    elif keep_overlap_only and keep_source == "indian":
        keep_class_names = overlap_names
    elif keep_overlap_only:
        keep_class_names = overlap_names
    elif keep_source == "spice_spectrum":
        keep_class_names = [c["name"] for c in uni["classes"]
                            if any(s.startswith("spice_spectrum:") for s in c["sources"])]
    elif keep_source == "indian":
        keep_class_names = [c["name"] for c in uni["classes"]
                            if any(s.startswith("indian:") for s in c["sources"])]
    else:
        keep_class_names = [c["name"] for c in uni["classes"]]

    # Re-number classes
    name_to_new_idx = {n: i for i, n in enumerate(sorted(keep_class_names))}
    old_idx_to_name = {c["index"]: c["name"] for c in uni["classes"]}

    # Filter samples
    for split in ("train", "val", "test"):
        for path, old_idx in uni["samples"][split]:
            name = old_idx_to_name[old_idx]
            if name not in name_to_new_idx:
                continue
            if keep_source is not None and _source_of(path) != keep_source:
                continue
            new_idx = name_to_new_idx[name]
            out["samples"][split].append([path, new_idx])

    # Rebuild class list with correct per-split counts after filtering
    name_to_counts = {n: {"train": 0, "val": 0, "test": 0} for n in keep_class_names}
    for split in ("train", "val", "test"):
        for _, idx in out["samples"][split]:
            inv = list(name_to_new_idx.keys())[idx]
            name_to_counts[inv][split] += 1
    for name in sorted(keep_class_names):
        c = next(c for c in uni["classes"] if c["name"] == name)
        cnt = name_to_counts[name]
        total = cnt["train"] + cnt["val"] + cnt["test"]
        out["classes"].append({
            "name": name, "index": name_to_new_idx[name],
            "sources": c["sources"], "total": total,
            "train": cnt["train"], "val": cnt["val"], "test": cnt["test"],
        })
    return out


def main():
    with open(UNIFIED) as f:
        uni = json.load(f)

    cfg = [
        ("manifest_ss_only.json",      {"keep_source": "spice_spectrum"}),
        ("manifest_indian_only.json",  {"keep_source": "indian"}),
        ("manifest_overlap_only.json", {"keep_overlap_only": True}),
        ("manifest_overlap_ss.json",   {"keep_source": "spice_spectrum", "keep_overlap_only": True}),
        ("manifest_overlap_indian.json",   {"keep_source": "indian", "keep_overlap_only": True}),
    ]
    for fname, kwargs in cfg:
        m = _filter_manifest(uni, **kwargs)
        out_path = OUT_DIR / fname
        with open(out_path, "w") as f:
            json.dump(m, f, indent=2)
        print(f"{fname}: {len(m['classes'])} classes | "
              f"train={len(m['samples']['train'])} "
              f"val={len(m['samples']['val'])} "
              f"test={len(m['samples']['test'])}")
        print(f"  classes: {[c['name'] for c in m['classes']]}")
        print(f"  saved: {out_path}\n")


if __name__ == "__main__":
    main()