#!/usr/bin/env python
"""Generate configs/wds_mix.yaml from converted wds dirs.

Scans <root> for dataset dirs (manifest.json at top level, or one level down
for bundles like unitree/<task>), weights sqrt(samples) — same rationale as
make_stream_specs.py: plain-proportional lets the biggest sets dominate.

Usage:
    python make_wds_mix.py --root data/wds --hf-prefix AlexWortega \
        > configs/wds_mix.yaml
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

# local dir name -> hf repo basename (bundles keep subdir)
REPO_FOR = {"unitree": "wds-unitree", "go_stanford": "wds-go-stanford",
            "recon": "wds-recon", "sacson": "wds-sacson",
            "scand": "wds-scand", "tartandrive": "wds-tartandrive"}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=Path, required=True)
    ap.add_argument("--hf-prefix", default="AlexWortega")
    args = ap.parse_args()

    rows = []
    for top in sorted(args.root.iterdir()):
        if not top.is_dir():
            continue
        repo = f"{args.hf_prefix}/{REPO_FOR.get(top.name, 'wds-' + top.name.replace('_', '-'))}"
        if (top / "manifest.json").exists():
            rows.append((repo, None, json.loads((top / "manifest.json").read_text())))
        else:
            for sub in sorted(top.iterdir()):
                if sub.is_dir() and (sub / "manifest.json").exists():
                    rows.append((repo, sub.name, json.loads((sub / "manifest.json").read_text())))

    total = sum(m["samples"] for _, _, m in rows)
    print(f"# autogenerated by make_wds_mix.py: {len(rows)} datasets, "
          f"{total:,} samples. weight = sqrt(samples).")
    print("datasets:")
    for repo, sub, m in rows:
        w = math.sqrt(m["samples"])
        print(f"  - hf_repo: {repo}")
        if sub:
            print(f"    subdir: {sub}")
        print(f"    weight: {w:.1f}   # {m['samples']:,} samples")
        print(f"    embodiment_id: {m['embodiment_id']}")


if __name__ == "__main__":
    main()