upload_to_hf.py

"""Upload CausalDriveBench to a Hugging Face dataset repo.

Stages the release-ready directory into a clean tree, sanitizes paths, renders
a dataset card from a template, and pushes everything via huggingface_hub.

Default flow:

  1. Copy --source into --staging/data/<subset>/ (preserves per-scene tree).
  2. Walk every frames.json under the staged subset; remove the `data_root`
     field. Per-scene files become location-agnostic.
  3. Build a sanitized top-level manifest.json from the source manifest. Strip
     `data_root` / `records` / `outputs` (internal paths) and add
     `nuscenes_root_placeholder` + a `splits` summary.
  4. Render README.md from the template + qa_statistics report.
  5. Copy LICENSE if present in scripts/release/.
  6. Upload --staging to HF via HfApi.upload_folder.

Auth: HF_TOKEN env var, or `huggingface-cli login` cached token.

Usage:

  python scripts/release/upload_to_hf.py --dry-run
  python scripts/release/upload_to_hf.py    # real upload
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import sys
from pathlib import Path

DEFAULT_SOURCE = Path(
    "/path/to/downloads/nuscenes_full_ice"
)
DEFAULT_STAGING = Path(
    "/path/to/downloads/causaldrivebench_release"
)
DEFAULT_REPO_ID = "causaldrivebench/CausalDriveBench"
ROOT_PLACEHOLDER = "${NUSCENES_ROOT}"

REPO_ROOT = Path(__file__).resolve().parents[2]
TEMPLATE_PATH = Path(__file__).with_name("dataset_card_template.md")
LICENSE_TEMPLATE = Path(__file__).with_name("LICENSE")
QA_STATS_DIR = REPO_ROOT / "misc"


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--source", type=Path, default=DEFAULT_SOURCE,
                   help="Release-ready dir to upload (default: nuscenes_full_ice).")
    p.add_argument("--staging", type=Path, default=DEFAULT_STAGING,
                   help="Local staging dir; mirrors HF repo structure.")
    p.add_argument("--repo-id", default=DEFAULT_REPO_ID,
                   help="HF dataset repo id (org/name).")
    p.add_argument("--subset", default="nuscenes",
                   choices=["nuscenes", "openscene", "argoverse2"],
                   help="Subset namespace under data/. Default nuscenes.")
    p.add_argument("--dry-run", action="store_true",
                   help="Stage + render locally; print upload plan; no network.")
    p.add_argument("--skip-stage", action="store_true",
                   help="Don't rebuild staging tree; just (re)upload what's there.")
    p.add_argument("--skip-card", action="store_true",
                   help="Don't regenerate README.md.")
    p.add_argument("--commit-message", default=None,
                   help="HF commit message (default auto-generated).")
    p.add_argument("--qa-stats-md", type=Path,
                   default=QA_STATS_DIR / "qa_statistics_nuscenes_full_ice_diverse.md",
                   help="QA stats markdown to embed in the dataset card. "
                        "Default is the 815-sample diverse subset (matches what we upload).")
    return p.parse_args()


# ─── Staging ────────────────────────────────────────────────────────────────

def stage_subset(source: Path, staging: Path, subset: str) -> Path:
    """Copy source into staging/data/<subset>/. Returns the subset dir."""
    if not source.is_dir():
        raise SystemExit(f"--source not found: {source}")
    target = staging / "data" / subset
    if target.exists():
        print(f"[stage] removing existing {target}")
        shutil.rmtree(target)
    target.parent.mkdir(parents=True, exist_ok=True)
    print(f"[stage] copy {source} -> {target}")
    shutil.copytree(source, target)
    # The source manifest.json gets copied inside the subset; we'll regenerate
    # a sanitized one at the staging root, so drop the inner one.
    inner_manifest = target / "manifest.json"
    if inner_manifest.exists():
        inner_manifest.unlink()
    return target


def sanitize_frames(subset_dir: Path) -> int:
    """Remove `data_root` from every per-scene frames.json. Returns count."""
    n = 0
    for fp in subset_dir.rglob("frames.json"):
        data = json.loads(fp.read_text())
        if "data_root" in data:
            data.pop("data_root")
            fp.write_text(json.dumps(data, indent=2))
            n += 1
    print(f"[sanitize] stripped data_root from {n} frames.json files")
    return n


# ─── QA flattener ───────────────────────────────────────────────────────────

# Common fields across all three QA flavours, plus per-flavour fields. Anything
# in this list ends up as a top-level column in qa.jsonl. `meta` and `options`
# are kept as nested JSON; HF's Parquet converter handles both.
_QA_COMMON = [
    "id", "rung", "category", "difficulty", "question", "answer_format",
    "correct_answer", "reasoning", "options",
]
_QA_PER_KIND = {
    "active":     ["graph_structure", "isolation_safe", "counterfactual_frame"],
    "dormant":    ["subtype", "spatial_context", "activation_plausibility"],
    "distractor": ["spatial_context"],
}
_QA_KINDS = {
    "active":     "active_qa.json",
    "dormant":    "dormant_qa.json",
    "distractor": "distractor_qa.json",
}


def _load_qa(path: Path) -> list:
    data = json.loads(path.read_text())
    if isinstance(data, dict) and "questions" in data:
        return data["questions"] or []
    if isinstance(data, list):
        return data
    return []


def build_qa_jsonl(subset_dir: Path, out_path: Path) -> int:
    """Walk the staged subset; emit one row per QA pair to out_path. Returns row count."""
    n = 0
    all_extras = sorted({k for v in _QA_PER_KIND.values() for k in v})
    with out_path.open("w") as out:
        for sample_dir in sorted(subset_dir.glob("*/SAMPLED_*")):
            scene_id = sample_dir.parent.name
            sample_id = sample_dir.name
            for kind, fname in _QA_KINDS.items():
                qpath = sample_dir / "qa" / fname
                if not qpath.exists():
                    continue
                for q in _load_qa(qpath):
                    row = {
                        "scene_id": scene_id,
                        "sample_id": sample_id,
                        "qa_type": kind,
                    }
                    for k in _QA_COMMON:
                        row[k] = q.get(k)
                    # Per-flavour fields: union of all extras, null when N/A
                    for k in all_extras:
                        row[k] = q.get(k)
                    # Stringify `meta` (nested + variable shape) for clean Parquet typing
                    meta = q.get("meta")
                    row["meta"] = json.dumps(meta) if meta is not None else None
                    out.write(json.dumps(row) + "\n")
                    n += 1
    print(f"[qa-jsonl] wrote {out_path} ({n:,} QA pairs)")
    return n


# ─── Manifest ───────────────────────────────────────────────────────────────

def build_manifest(source: Path, subset: str) -> dict:
    src_manifest_path = source / "manifest.json"
    if not src_manifest_path.exists():
        raise SystemExit(f"source manifest not found: {src_manifest_path}")
    src = json.loads(src_manifest_path.read_text())

    # Distinct scene count from the samples list.
    samples = src.get("samples", []) or []
    scenes = sorted({s.get("scene") for s in samples if s.get("scene")})

    out = {
        "dataset_name": "causaldrivebench",
        "version": "v1",
        "nuscenes_root_placeholder": ROOT_PLACEHOLDER,
        "n_samples": src.get("n_samples", len(samples)),
        "splits": {
            subset: {
                "n_samples": len(samples),
                "n_scenes": len(scenes),
            }
        },
        "filtered_to": src.get("filtered_to"),
        "missing_graph": src.get("missing_graph", []),
        "missing_qa": src.get("missing_qa", []),
        "samples": samples,
    }
    # Drop internal absolute paths; keep nothing that leaks our infra.
    for stripped in ("data_root", "records", "outputs"):
        out.pop(stripped, None)
    return out


def write_manifest(staging: Path, manifest: dict) -> Path:
    p = staging / "manifest.json"
    p.write_text(json.dumps(manifest, indent=2))
    print(f"[manifest] wrote {p} (n_samples={manifest['n_samples']})")
    return p


# ─── Dataset card ───────────────────────────────────────────────────────────

_INTERNAL_PATH_RE = re.compile(r"/spiral_hdd_\d+\S*")


def _sanitize_embed(text: str) -> str:
    """Strip lines that reference absolute paths on our infra."""
    out_lines = []
    for line in text.splitlines():
        if _INTERNAL_PATH_RE.search(line):
            continue
        out_lines.append(line)
    return "\n".join(out_lines)


def render_card(
    template_path: Path,
    manifest: dict,
    qa_stats_md: Path | None,
    subset: str,
) -> str:
    if not template_path.exists():
        raise SystemExit(f"template not found: {template_path}")
    tpl = template_path.read_text()

    qa_stats_block = ""
    if qa_stats_md and qa_stats_md.exists():
        qa_stats_block = _sanitize_embed(qa_stats_md.read_text().strip())
    else:
        qa_stats_block = "_QA stats report not found at upload time._"

    split = manifest["splits"][subset]
    fields = {
        "{{SUBSET}}": subset,
        "{{N_SAMPLES}}": str(split["n_samples"]),
        "{{N_SCENES}}": str(split["n_scenes"]),
        "{{ROOT_PLACEHOLDER}}": ROOT_PLACEHOLDER,
        "{{QA_STATS}}": qa_stats_block,
    }
    for k, v in fields.items():
        tpl = tpl.replace(k, v)
    return tpl


def write_card(staging: Path, body: str) -> Path:
    p = staging / "README.md"
    p.write_text(body)
    print(f"[card] wrote {p}")
    return p


# ─── License ────────────────────────────────────────────────────────────────

def copy_license(staging: Path) -> None:
    if not LICENSE_TEMPLATE.exists():
        print(f"[license] WARNING: template missing at {LICENSE_TEMPLATE}; "
              "skipping. Add one before the real upload.")
        return
    dst = staging / "LICENSE"
    shutil.copy(LICENSE_TEMPLATE, dst)
    print(f"[license] copied {LICENSE_TEMPLATE} -> {dst}")


# ─── Upload plan / push ─────────────────────────────────────────────────────

def summarize_plan(staging: Path) -> tuple[int, int]:
    """Walk staging dir; return (file_count, total_bytes)."""
    files = [p for p in staging.rglob("*") if p.is_file()]
    total = sum(p.stat().st_size for p in files)
    return len(files), total


def fmt_bytes(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024:
            return f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} TB"


def push_to_hf(
    staging: Path,
    repo_id: str,
    commit_message: str | None,
) -> None:
    try:
        from huggingface_hub import HfApi, create_repo
    except ImportError:
        raise SystemExit(
            "huggingface_hub not installed. Run:\n"
            "  pip install huggingface_hub\n"
            "or `conda activate causal && pip install huggingface_hub`."
        )

    token = os.environ.get("HF_TOKEN")
    api = HfApi(token=token)

    try:
        api.whoami()
    except Exception as e:
        raise SystemExit(
            f"HF auth failed: {e}\n"
            "Set HF_TOKEN env var or run `huggingface-cli login`."
        )

    create_repo(repo_id, repo_type="dataset", exist_ok=True, token=token)

    msg = commit_message or "upload causaldrivebench v1 (nuscenes subset)"
    print(f"[upload] pushing {staging} -> hf://datasets/{repo_id}")
    api.upload_folder(
        folder_path=str(staging),
        repo_id=repo_id,
        repo_type="dataset",
        commit_message=msg,
        # croissant.json is fetched/managed separately via fetch_croissant.py;
        # don't ship a stale local copy with the data upload.
        ignore_patterns=[".DS_Store", "*.pyc", "__pycache__/*", "croissant.json"],
    )

    files = api.list_repo_files(repo_id, repo_type="dataset")
    print(f"[upload] OK. {len(files)} files now on the repo.")


# ─── Main ───────────────────────────────────────────────────────────────────

def main() -> None:
    args = parse_args()
    staging: Path = args.staging.resolve()
    source: Path = args.source.resolve()

    if not args.skip_stage:
        subset_dir = stage_subset(source, staging, args.subset)
        sanitize_frames(subset_dir)
    else:
        subset_dir = staging / "data" / args.subset
        if not subset_dir.is_dir():
            raise SystemExit(
                f"--skip-stage but {subset_dir} doesn't exist. "
                "Run without --skip-stage at least once first."
            )

    build_qa_jsonl(subset_dir, staging / "qa.jsonl")

    manifest = build_manifest(source, args.subset)
    write_manifest(staging, manifest)

    if not args.skip_card:
        body = render_card(TEMPLATE_PATH, manifest, args.qa_stats_md, args.subset)
        write_card(staging, body)

    copy_license(staging)

    n_files, total = summarize_plan(staging)
    print(
        f"[plan] {n_files} files, {fmt_bytes(total)} -> "
        f"hf://datasets/{args.repo_id}"
    )

    if args.dry_run:
        print("[dry-run] skipping upload. Inspect:")
        print(f"  {staging}/manifest.json")
        print(f"  {staging}/README.md")
        print(f"  {staging}/data/{args.subset}/<scene>/<sample>/frames.json")
        return

    push_to_hf(staging, args.repo_id, args.commit_message)


if __name__ == "__main__":
    main()