"""
make_loso_manifests.py — Leave-One-Source-Out manifests for domain generalization.

For each acquisition source S present in the base manifest, writes a manifest that
TRAINS on all *other* sources and TESTS on S (held out entirely), restricted to the
classes shared between S and the training sources, reindexed 0..k-1.

This is the gold-standard DG protocol. With the current 2 real sources it produces
the two cross-source splits; it scales automatically as more sources are added
(n>=3 is where AIFNet's adversary/CORAL become a genuine DG contribution).

Output: outputs/manifest_loso_<source>.json  (same schema as unified_benchmark.json,
so src.dataset.load_manifest_splits / get_aifnet_dataloaders consume it directly).

    python make_loso_manifests.py                          # from unified_benchmark.json
    python make_loso_manifests.py --base outputs/unified_benchmark.json --val-frac 0.15
"""
import argparse
import json
from collections import defaultdict
from pathlib import Path

import numpy as np


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


def load_samples(base):
    m = json.load(open(base))
    idx2name = {c["index"]: c["name"] for c in m["classes"]}
    # (path, class_name, source) pooled across the base splits
    samples = []
    for split in ("train", "val", "test"):
        for path, label in m["samples"][split]:
            samples.append((path, idx2name[int(label)], source_of(path)))
    return samples


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="outputs/unified_benchmark.json")
    ap.add_argument("--val-frac", type=float, default=0.15)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--outdir", default="outputs")
    args = ap.parse_args()

    samples = load_samples(args.base)
    sources = sorted({s for _, _, s in samples})
    by_src_cls = defaultdict(lambda: defaultdict(list))
    for p, c, s in samples:
        by_src_cls[s][c].append(p)
    print(f"sources={sources}  classes/source=" +
          ", ".join(f"{s}:{len(by_src_cls[s])}" for s in sources))

    if len(sources) < 2:
        raise SystemExit("need >=2 sources for LOSO")

    written = []
    for held in sources:
        train_srcs = [s for s in sources if s != held]
        train_classes = set().union(*[set(by_src_cls[s]) for s in train_srcs])
        common = sorted(set(by_src_cls[held]) & train_classes)
        if not common:
            print(f"  [skip] {held}: no shared classes with training sources")
            continue
        cls2idx = {c: i for i, c in enumerate(common)}
        rng = np.random.default_rng(args.seed)

        train, val, test = [], [], []
        for s in train_srcs:                       # train/val from the other sources
            for c in common:
                paths = list(by_src_cls[s][c])
                rng.shuffle(paths)
                n_val = max(1, int(round(len(paths) * args.val_frac))) if paths else 0
                for p in paths[:n_val]:
                    val.append([p, cls2idx[c]])
                for p in paths[n_val:]:
                    train.append([p, cls2idx[c]])
        for c in common:                           # test = the entire held-out source
            for p in by_src_cls[held][c]:
                test.append([p, cls2idx[c]])

        manifest = {
            "version": 1, "seed": args.seed, "protocol": "leave-one-source-out",
            "held_out_source": held, "train_sources": train_srcs,
            "classes": [{"name": c, "index": cls2idx[c]} for c in common],
            "splits": {"train": 1 - args.val_frac, "val": args.val_frac, "test": "held-out source"},
            "samples": {"train": train, "val": val, "test": test},
        }
        out = Path(args.outdir) / f"manifest_loso_{held}.json"
        json.dump(manifest, open(out, "w"))
        written.append(out)
        print(f"  {out.name}: {len(common)} classes | train {len(train)} val {len(val)} "
              f"test {len(test)} (held-out {held})")

    print(f"\nwrote {len(written)} LOSO manifests. With n={len(sources)} sources, "
          f"each trains on {len(sources)-1} and tests on 1.")


if __name__ == "__main__":
    main()