"""Spatial-context audit across representative LeRobot datasets.

Reproduces the 2026-07-23 audit behind
docs/blog-post-spatial-audit-2026-07-23.md. For each dataset, pinned to
the revision SHA observed at audit time (HF main branches are mutable):

1. `traceplane check` (metadata-level, `check_data=False`) — format,
   episode counts, and the spatial-context verdict (which looks for
   Traceplane's versioned spatial artifact: meta/spatial_context.json +
   spatial/relations.parquet).
2. A feature-declaration scan of `meta/info.json` — which observation
   streams the dataset *declares*. This is a declaration check, not a
   content check: it cannot see sidecar files, dataset-card docs, or
   unconventionally named features.

Usage: python scripts/run-audit-spatial.py [out.json]
Requires: pip install traceplane>=0.5.9; anonymous HF access (the audit
ran with no HF token; heavy re-runs may hit anonymous rate limits).
"""
import gzip
import json
import sys

import fsspec

from traceplane.check import check_dataset

# Revision SHAs recorded 2026-07-23 via the x-repo-commit response
# header of resolve/main/meta/info.json.
DATASETS = {
    "lerobot/pusht": "7628202a2180972f291ba1bc6723834921e72c19",
    "lerobot/aloha_sim_transfer_cube_human": "6a43d500f101255823a9d2b9dc244eeb01a2cd31",
    "lerobot/droid_100": "87301a2d2e99340e2010c9ef0f1d8e780b08aaf9",
    "lerobot/unitreeh1_fold_clothes": "a0926b4dd672560fdc4e6c9d1e68cda17631d608",
    "lerobot/svla_so101_pickplace": "f641879e22172be7e8161d5e6c1503c2d2feb657",
    "unitreerobotics/G1_Dex3_GraspSquare_Dataset": "c2b44ca4c5f6329d75143613f9f1eada0cf94c62",
    # USC-GVL/humanoid-everyday redirects to USC-PSI-Lab (repo renamed).
    "USC-GVL/humanoid-everyday": "71f6210e91c2a7e8b8e4d47db0147444a1a0640e",
}

CALIB_HINTS = ("intrinsic", "extrinsic", "calib", "camera_param", "cam_ext", "cam_int", "transform")
DEPTH_HINTS = ("depth",)


def scan_declared_features(fs, path: str) -> dict:
    with fs.open(f"{path}/meta/info.json", "rb") as h:
        raw = h.read()
    if raw[:2] == b"\x1f\x8b":  # gzip-compressed on some HF repos
        raw = gzip.decompress(raw)
    info = json.loads(raw)
    names = list(info.get("features") or {})
    return {
        "n_features": len(names),
        "declared_depth": [n for n in names if any(k in n.lower() for k in DEPTH_HINTS)],
        "declared_calibration": [n for n in names if any(k in n.lower() for k in CALIB_HINTS)],
        "cameras": [n for n in names if "image" in n.lower() or "rgb" in n.lower()],
        "robot_type": info.get("robot_type"),
    }


def main() -> None:
    results = []
    for name, revision in DATASETS.items():
        uri = f"hf://datasets/{name}@{revision}"
        row = {"dataset": name, "revision": revision}
        try:
            report = check_dataset(uri, check_data=False, storage_options={"token": False})
            row["format"] = report.format_version or "unrecognized"
            row["episodes"] = report.total_episodes
            row["versioned_spatial_artifact"] = report.spatial_context or "absent"
            row["errors"] = len(report.errors)
            row["warnings"] = len(report.warnings)
            fs, path = fsspec.core.url_to_fs(uri, token=False)
            row.update(scan_declared_features(fs, path))
        except Exception as exc:  # noqa: BLE001
            row["error"] = f"{type(exc).__name__}: {exc}"
        results.append(row)
        print(json.dumps(row), flush=True)

    out = sys.argv[1] if len(sys.argv) > 1 else "spatial-audit-results.json"
    with open(out, "w") as f:
        json.dump(results, f, indent=2)
    print(f"wrote {out}")


if __name__ == "__main__":
    main()