#!/usr/bin/env python
"""Probe the nvidia/physical-ai collection and emit a streaming specs YAML.

Walks the robotics repos, finds every LeRobot v2.x (sub)dataset that
HubEpisodeStream can consume, assigns embodiment ids BY ROBOT FAMILY (not per
dataset — 127 ids would defeat the embedding) and weights proportional to
sqrt(frames) (plain proportional lets the two biggest GR1 tasks dominate the
mixture; sqrt keeps diversity without starving the big sources).

v3.0 repos (BridgeData2_LeRobot_v3, LIBERO_LeRobot_v3) are NOT emitted — they
need the v3 shard path (build_shards.py); both were already in the C-scaled
pretraining mixture anyway.

Usage:
    python make_stream_specs.py --out configs/physical_ai_stream.yaml
"""

from __future__ import annotations

import argparse
import json
import math
import os
import urllib.request
from concurrent.futures import ThreadPoolExecutor

FAMILIES = [  # ordered: first match wins; ids continue after C-scaled's 0..9
    ("gr1", 10), ("bimanual_panda_gripper", 11), ("bimanual_panda_hand", 12),
    ("single_panda_gripper", 13), ("sim_behavior_r1_pro", 14), ("unitree_g1", 15),
    ("g1", 15),
]

ROOTS = [
    ("nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim", None),   # nested
    ("nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim", "LeRobot"),    # nested under LeRobot/
    ("nvidia/PhysicalAI-GR00T-Tuned-Tasks", None),                 # nested
    ("nvidia/PhysicalAI-Robotics-GR00T-Teleop-G1", None),          # nested
    ("nvidia/GR00T-N1.7-AppleToPlate", ""),                        # root-level v2.1
]


def hf_json(url: str, token: str | None):
    req = urllib.request.Request(url, headers={"Authorization": f"Bearer {token}"} if token else {})
    try:
        return json.load(urllib.request.urlopen(req, timeout=60))
    except Exception:
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="configs/physical_ai_stream.yaml")
    ap.add_argument("--hz", type=float, default=10.0)
    args = ap.parse_args()
    token = None
    tp = os.path.expanduser("~/.cache/huggingface/token")
    if os.path.exists(tp):
        token = open(tp).read().strip()

    jobs = []  # (repo_id, prefix)
    for repo, sub in ROOTS:
        if sub == "":  # root-level dataset
            jobs.append((repo, ""))
            continue
        base = f"https://huggingface.co/api/datasets/{repo}/tree/main" + (f"/{sub}" if sub else "")
        tree = hf_json(base, token) or []
        for e in tree:
            if e.get("type") == "directory":
                jobs.append((repo, e["path"]))

    def probe(job):
        repo, pfx = job
        p = f"{pfx}/" if pfx else ""
        info = hf_json(f"https://huggingface.co/datasets/{repo}/resolve/main/{p}meta/info.json", token)
        if not info or "total_frames" not in info:
            return None
        ver = str(info.get("codebase_version", ""))
        if not ver.startswith("v2"):
            return None  # v3 needs the shard path, v1 untested
        return dict(
            repo_id=repo, prefix=pfx, frames=info["total_frames"], fps=info["fps"],
            state_dim=info["features"]["observation.state"]["shape"][0],
            action_dim=info["features"]["action"]["shape"][0],
            ncam=sum(1 for v in info["features"].values() if v.get("dtype") == "video"),
        )

    with ThreadPoolExecutor(16) as ex:
        found = [r for r in ex.map(probe, jobs) if r]

    max_state = max(r["state_dim"] for r in found)
    max_action = max(r["action_dim"] for r in found)
    total = sum(r["frames"] for r in found)
    lines = [
        "# Autogenerated by make_stream_specs.py — streaming mixture over the",
        "# LeRobot-v2.x part of the nvidia/physical-ai collection.",
        f"# {len(found)} datasets, {total:,} source frames.",
        f"# max_state_dim needed: {max_state}; max_action_dim needed: {max_action}.",
        "datasets:",
    ]
    for r in sorted(found, key=lambda x: (-x["frames"])):
        name = r["prefix"] or r["repo_id"].split("/")[-1]
        emb = next((i for k, i in FAMILIES if name.startswith(k) or f"/{k}" in name.lower() or k in name.lower()), 15)
        w = math.sqrt(r["frames"])
        lines += [
            f"  - repo_id: {r['repo_id']}",
            f"    prefix: \"{r['prefix']}\"",
            f"    weight: {w:.1f}   # sqrt({r['frames']:,} fr); fps={r['fps']} state={r['state_dim']} act={r['action_dim']} cams={r['ncam']}",
            f"    embodiment_id: {emb}",
            f"    hz: {args.hz}",
        ]
    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    with open(args.out, "w") as f:
        f.write("\n".join(lines) + "\n")
    print(f"{len(found)} датасетов, {total:,} кадров -> {args.out}")
    print(f"max_state_dim={max_state} max_action_dim={max_action}")


if __name__ == "__main__":
    main()