scripts/data/rebuild_trex_dataset_variants.py

#!/usr/bin/env python3
"""Rebuild filtered T-Rex LeRobot-v2 base/Track-Force dataset variants.

The builder is resumable and writes only under a staging directory until the
explicit install phase. RGB/tactile videos are hard-linked; Parquet files and
renumbered Track NPZ files are rewritten atomically.
"""

from __future__ import annotations

import argparse
import copy
import hashlib
import json
import math
import os
import shutil
import sys
import tempfile
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable, Sequence

import numpy as np

_DREAMZERO_ROOT = Path(__file__).resolve().parents[2]
_SCRIPTS_ROOT = _DREAMZERO_ROOT / "scripts"
for _path in (_SCRIPTS_ROOT, _SCRIPTS_ROOT / "data"):
    if str(_path) not in sys.path:
        sys.path.insert(0, str(_path))

import build_trex_track_force_v2 as force_builder  # noqa: E402

TRACK_CACHE_NAME = "tracks_trex_track_force_v2"
VARIANT_NAMES = (
    "trex_small",
    "trex_small_force",
    "trex_full",
    "trex_full_force",
)
LEGACY_NAMES = ("trex_small", "trex_datasetv2", "trex_track_force_v2")
PARQUET_SCHEMA_METADATA_KEY = b"trex_track_force_schema_version"


@dataclass(frozen=True)
class EpisodeRecord:
    source_episode_index: int
    episode_index: int
    length: int
    task: str
    task_index: int
    global_index_start: int


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _read_json(path: Path) -> dict:
    return json.loads(path.read_text())


def _read_jsonl(path: Path) -> list[dict]:
    return [
        json.loads(line)
        for line in path.read_text().splitlines()
        if line.strip()
    ]


def _atomic_write_json(path: Path, value: object) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary_name = tempfile.mkstemp(
        prefix=f".{path.name}.",
        suffix=".tmp",
        dir=path.parent,
    )
    try:
        with os.fdopen(fd, "w") as file:
            json.dump(value, file, indent=2)
            file.write("\n")
        os.replace(temporary_name, path)
    except Exception:
        Path(temporary_name).unlink(missing_ok=True)
        raise


def _atomic_write_jsonl(path: Path, rows: Iterable[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary_name = tempfile.mkstemp(
        prefix=f".{path.name}.",
        suffix=".tmp",
        dir=path.parent,
    )
    try:
        with os.fdopen(fd, "w") as file:
            for row in rows:
                file.write(json.dumps(row) + "\n")
        os.replace(temporary_name, path)
    except Exception:
        Path(temporary_name).unlink(missing_ok=True)
        raise


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as file:
        while chunk := file.read(8 * 1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def _video_keys(info: dict) -> list[str]:
    return sorted(
        key
        for key, feature in info["features"].items()
        if feature.get("dtype") == "video"
    )


def _source_episode_path(root: Path, episode_index: int, suffix: str) -> Path:
    return (
        root
        / suffix
        / f"chunk-{episode_index // 1000:03d}"
        / f"episode_{episode_index:06d}.parquet"
    )


def _episode_parquet_path(root: Path, episode_index: int) -> Path:
    return (
        root
        / "data"
        / f"chunk-{episode_index // 1000:03d}"
        / f"episode_{episode_index:06d}.parquet"
    )


def _episode_video_path(root: Path, episode_index: int, video_key: str) -> Path:
    return (
        root
        / "videos"
        / f"chunk-{episode_index // 1000:03d}"
        / video_key
        / f"episode_{episode_index:06d}.mp4"
    )


def _track_path(root: Path, episode_index: int) -> Path:
    return root / TRACK_CACHE_NAME / f"episode_{episode_index:06d}.npz"


def _build_registry(
    episodes: Sequence[dict],
    source_tasks: Sequence[dict],
    *,
    excluded: set[int],
    source_limit: int | None,
) -> tuple[list[EpisodeRecord], list[dict]]:
    tasks_by_text = {
        str(row["task"]): int(row["task_index"])
        for row in source_tasks
    }
    selected = [
        episode
        for episode in episodes
        if int(episode["episode_index"]) not in excluded
        and (
            source_limit is None
            or int(episode["episode_index"]) < int(source_limit)
        )
    ]
    retained_tasks = {
        str(task)
        for episode in selected
        for task in episode["tasks"]
    }
    missing_tasks = retained_tasks.difference(tasks_by_text)
    if missing_tasks:
        raise ValueError(f"episodes reference unknown tasks: {sorted(missing_tasks)[:3]}")
    ordered_tasks = sorted(retained_tasks, key=tasks_by_text.__getitem__)
    new_task_by_text = {
        task: task_index
        for task_index, task in enumerate(ordered_tasks)
    }
    task_rows = [
        {"task_index": task_index, "task": task}
        for task, task_index in new_task_by_text.items()
    ]

    registry: list[EpisodeRecord] = []
    global_index_start = 0
    for new_episode_index, episode in enumerate(selected):
        tasks = [str(task) for task in episode["tasks"]]
        if len(tasks) != 1:
            raise ValueError(
                f"episode {episode['episode_index']} has {len(tasks)} tasks; expected one"
            )
        length = int(episode["length"])
        registry.append(
            EpisodeRecord(
                source_episode_index=int(episode["episode_index"]),
                episode_index=new_episode_index,
                length=length,
                task=tasks[0],
                task_index=new_task_by_text[tasks[0]],
                global_index_start=global_index_start,
            )
        )
        global_index_start += length
    return registry, task_rows


def _replace_primitive_column(table, name: str, values: np.ndarray):
    import pyarrow as pa

    index = table.schema.get_field_index(name)
    if index < 0:
        raise ValueError(f"missing required parquet column {name!r}")
    field = table.schema.field(index)
    return table.set_column(index, field, pa.array(values, type=field.type))


def _parquet_is_ready(
    path: Path,
    *,
    episode_index: int,
    task_index: int,
    global_index_start: int,
    expected_rows: int,
    force: bool,
) -> bool:
    if not path.is_file():
        return False
    try:
        import pyarrow.parquet as pq

        parquet_file = pq.ParquetFile(path)
        if int(parquet_file.metadata.num_rows) != int(expected_rows):
            return False
        metadata = parquet_file.schema_arrow.metadata or {}
        if force and metadata.get(PARQUET_SCHEMA_METADATA_KEY) != (
            force_builder.SCHEMA_VERSION.encode()
        ):
            return False
        table = pq.read_table(
            path,
            columns=["episode_index", "task_index", "frame_index", "index"],
        )
        stored_episode = np.asarray(table["episode_index"].to_numpy())
        stored_task = np.asarray(table["task_index"].to_numpy())
        frame_index = np.asarray(table["frame_index"].to_numpy())
        global_index = np.asarray(table["index"].to_numpy())
        return bool(
            np.all(stored_episode == int(episode_index))
            and np.all(stored_task == int(task_index))
            and np.array_equal(frame_index, np.arange(expected_rows))
            and np.array_equal(
                global_index,
                np.arange(
                    global_index_start,
                    global_index_start + expected_rows,
                ),
            )
        )
    except Exception:
        return False


def _rewrite_parquet_job(job: tuple) -> tuple[int, str]:
    (
        source_path_text,
        destination_path_text,
        episode_index,
        task_index,
        global_index_start,
        expected_rows,
        force,
    ) = job
    source_path = Path(source_path_text)
    destination_path = Path(destination_path_text)
    if _parquet_is_ready(
        destination_path,
        episode_index=episode_index,
        task_index=task_index,
        global_index_start=global_index_start,
        expected_rows=expected_rows,
        force=force,
    ):
        return episode_index, "skipped"

    import pyarrow.parquet as pq

    table = pq.read_table(source_path)
    if int(table.num_rows) != int(expected_rows):
        raise ValueError(
            f"{source_path}: {table.num_rows} rows != expected {expected_rows}"
        )
    original_metadata = table.schema.metadata
    table = _replace_primitive_column(
        table,
        "episode_index",
        np.full(expected_rows, episode_index, dtype=np.int64),
    )
    table = _replace_primitive_column(
        table,
        "task_index",
        np.full(expected_rows, task_index, dtype=np.int64),
    )
    table = _replace_primitive_column(
        table,
        "frame_index",
        np.arange(expected_rows, dtype=np.int64),
    )
    table = _replace_primitive_column(
        table,
        "index",
        np.arange(
            global_index_start,
            global_index_start + expected_rows,
            dtype=np.int64,
        ),
    )
    table = table.replace_schema_metadata(original_metadata)
    if force:
        metadata = table.schema.metadata or {}
        if metadata.get(PARQUET_SCHEMA_METADATA_KEY) != (
            force_builder.SCHEMA_VERSION.encode()
        ):
            raise ValueError(f"{source_path}: force schema metadata is missing")

    destination_path.parent.mkdir(parents=True, exist_ok=True)
    temporary_path = destination_path.with_suffix(".tmp.parquet")
    temporary_path.unlink(missing_ok=True)
    try:
        pq.write_table(
            table,
            temporary_path,
            compression="zstd",
            use_dictionary=True,
        )
        os.replace(temporary_path, destination_path)
    except Exception:
        temporary_path.unlink(missing_ok=True)
        raise
    return episode_index, "written"


def _track_is_ready(
    path: Path,
    *,
    episode_index: int,
    expected_frames: int,
) -> bool:
    if not path.is_file():
        return False
    try:
        with np.load(path, allow_pickle=False) as payload:
            return bool(
                int(np.asarray(payload["episode_index"]).item())
                == int(episode_index)
                and int(np.asarray(payload["num_steps"]).item())
                == int(expected_frames)
            )
    except Exception:
        return False


def _rewrite_track_job(job: tuple) -> tuple[int, str, str]:
    (
        source_path_text,
        destination_path_text,
        source_episode_index,
        episode_index,
        expected_frames,
    ) = job
    source_path = Path(source_path_text)
    destination_path = Path(destination_path_text)
    if _track_is_ready(
        destination_path,
        episode_index=episode_index,
        expected_frames=expected_frames,
    ):
        return episode_index, _sha256(destination_path), "skipped"

    destination_path.parent.mkdir(parents=True, exist_ok=True)
    if int(source_episode_index) == int(episode_index):
        destination_path.unlink(missing_ok=True)
        os.link(source_path, destination_path)
        return episode_index, _sha256(destination_path), "linked"

    with np.load(source_path, allow_pickle=False) as archive:
        payload = {
            name: np.asarray(archive[name]).copy()
            for name in archive.files
        }
    payload["episode_index"] = np.array(episode_index, dtype=np.int32)
    payload["source_episode_index"] = np.array(
        source_episode_index,
        dtype=np.int32,
    )
    temporary_path = destination_path.with_suffix(".tmp.npz")
    temporary_path.unlink(missing_ok=True)
    try:
        with temporary_path.open("wb") as file:
            np.savez_compressed(file, **payload)
        os.replace(temporary_path, destination_path)
    except Exception:
        temporary_path.unlink(missing_ok=True)
        raise
    return episode_index, _sha256(destination_path), "written"


def _hardlink(source: Path, destination: Path) -> str:
    if destination.exists():
        if os.path.samefile(source, destination):
            return "skipped"
        raise FileExistsError(f"destination is not the expected hard link: {destination}")
    destination.parent.mkdir(parents=True, exist_ok=True)
    os.link(source, destination)
    return "linked"


def _link_job(job: tuple[str, str]) -> str:
    return _hardlink(Path(job[0]), Path(job[1]))


def _run_jobs(
    jobs: Sequence[tuple],
    worker,
    *,
    workers: int,
    process: bool,
    label: str,
) -> list:
    if not jobs:
        return []
    executor_type = ProcessPoolExecutor if process else ThreadPoolExecutor
    results = []
    with executor_type(max_workers=max(1, workers)) as executor:
        futures = [executor.submit(worker, job) for job in jobs]
        for completed, future in enumerate(as_completed(futures), start=1):
            results.append(future.result())
            if completed % 250 == 0 or completed == len(futures):
                print(f"{label}: {completed}/{len(futures)}", flush=True)
    return results


def _write_variant_metadata_skeleton(
    *,
    variant_root: Path,
    final_root: Path,
    source_info: dict,
    source_modality: dict,
    source_embodiment: dict,
    registry: Sequence[EpisodeRecord],
    tasks: Sequence[dict],
    video_keys: Sequence[str],
    excluded_source_indices: Sequence[int],
    force: bool,
    source_dataset: Path,
    blacklist_sha256: str,
) -> None:
    info = copy.deepcopy(source_info)
    info["total_episodes"] = len(registry)
    info["total_frames"] = int(sum(record.length for record in registry))
    info["total_tasks"] = len(tasks)
    info["total_videos"] = len(registry) * len(video_keys)
    info["total_chunks"] = math.ceil(len(registry) / int(info["chunks_size"]))
    info["splits"] = {"train": f"0:{len(registry)}"}
    info.pop("discarded_episode_indices", None)
    info.pop("trex_track_force", None)
    info["trex_filter"] = {
        "source_dataset": str(source_dataset),
        "source_episode_count": int(source_info["total_episodes"]),
        "source_to_new_map": "meta/source_episode_index_map.json",
        "excluded_source_indices": "meta/excluded_source_episode_indices.json",
        "blacklist_sha256": blacklist_sha256,
        "force_variant": bool(force),
        "created_at": _utc_now(),
    }
    metadata_dir = variant_root / "meta"
    metadata_dir.mkdir(parents=True, exist_ok=True)
    _atomic_write_json(metadata_dir / "info.json", info)
    _atomic_write_json(metadata_dir / "modality.json", source_modality)
    _atomic_write_json(metadata_dir / "embodiment.json", source_embodiment)
    _atomic_write_jsonl(
        metadata_dir / "episodes.jsonl",
        (
            {
                "episode_index": record.episode_index,
                "tasks": [record.task],
                "length": record.length,
            }
            for record in registry
        ),
    )
    _atomic_write_jsonl(metadata_dir / "tasks.jsonl", tasks)
    _atomic_write_json(
        metadata_dir / "source_episode_index_map.json",
        {
            str(record.source_episode_index): record.episode_index
            for record in registry
        },
    )
    _atomic_write_json(
        metadata_dir / "episode_index_provenance.json",
        [asdict(record) for record in registry],
    )
    _atomic_write_json(
        metadata_dir / "excluded_source_episode_indices.json",
        list(excluded_source_indices),
    )
    _atomic_write_json(
        metadata_dir / "dataset_variant.json",
        {
            "name": final_root.name,
            "final_root": str(final_root),
            "force": bool(force),
            "episodes": len(registry),
            "frames": int(sum(record.length for record in registry)),
            "tasks": len(tasks),
            "video_keys": list(video_keys),
            "created_at": _utc_now(),
        },
    )


def _rewrite_parquets(
    *,
    source_root: Path,
    destination_root: Path,
    registry: Sequence[EpisodeRecord],
    workers: int,
    force: bool,
) -> None:
    jobs = [
        (
            str(_episode_parquet_path(source_root, record.source_episode_index)),
            str(_episode_parquet_path(destination_root, record.episode_index)),
            record.episode_index,
            record.task_index,
            record.global_index_start,
            record.length,
            force,
        )
        for record in registry
    ]
    _run_jobs(
        jobs,
        _rewrite_parquet_job,
        workers=workers,
        process=True,
        label=f"{destination_root.name} parquet",
    )


def _link_videos(
    *,
    source_root: Path,
    destination_root: Path,
    registry: Sequence[EpisodeRecord],
    video_keys: Sequence[str],
    source_uses_new_indices: bool,
    workers: int,
) -> None:
    jobs: list[tuple[str, str]] = []
    for record in registry:
        source_episode_index = (
            record.episode_index
            if source_uses_new_indices
            else record.source_episode_index
        )
        for video_key in video_keys:
            jobs.append(
                (
                    str(
                        _episode_video_path(
                            source_root,
                            source_episode_index,
                            video_key,
                        )
                    ),
                    str(
                        _episode_video_path(
                            destination_root,
                            record.episode_index,
                            video_key,
                        )
                    ),
                )
            )
    _run_jobs(
        jobs,
        _link_job,
        workers=workers,
        process=False,
        label=f"{destination_root.name} videos",
    )


def _rewrite_tracks(
    *,
    source_root: Path,
    destination_root: Path,
    registry: Sequence[EpisodeRecord],
    workers: int,
) -> dict[int, str]:
    jobs = [
        (
            str(_track_path(source_root, record.source_episode_index)),
            str(_track_path(destination_root, record.episode_index)),
            record.source_episode_index,
            record.episode_index,
            record.length,
        )
        for record in registry
    ]
    results = _run_jobs(
        jobs,
        _rewrite_track_job,
        workers=workers,
        process=True,
        label=f"{destination_root.name} tracks",
    )
    return {int(episode_index): sha256 for episode_index, sha256, _ in results}


def _link_subset_files(
    *,
    source_root: Path,
    destination_root: Path,
    registry: Sequence[EpisodeRecord],
    video_keys: Sequence[str],
    force: bool,
    workers: int,
) -> dict[int, str]:
    parquet_jobs = [
        (
            str(_episode_parquet_path(source_root, record.episode_index)),
            str(_episode_parquet_path(destination_root, record.episode_index)),
        )
        for record in registry
    ]
    _run_jobs(
        parquet_jobs,
        _link_job,
        workers=workers,
        process=False,
        label=f"{destination_root.name} parquet links",
    )
    _link_videos(
        source_root=source_root,
        destination_root=destination_root,
        registry=registry,
        video_keys=video_keys,
        source_uses_new_indices=True,
        workers=workers,
    )
    track_hashes: dict[int, str] = {}
    if force:
        track_jobs = [
            (
                str(_track_path(source_root, record.episode_index)),
                str(_track_path(destination_root, record.episode_index)),
            )
            for record in registry
        ]
        _run_jobs(
            track_jobs,
            _link_job,
            workers=workers,
            process=False,
            label=f"{destination_root.name} track links",
        )
        track_hashes = {
            record.episode_index: _sha256(
                _track_path(destination_root, record.episode_index)
            )
            for record in registry
        }
    return track_hashes


def _aggregate_source_manifest_entries(source_force_root: Path) -> dict[int, dict]:
    manifests = sorted(
        (
            source_force_root
            / "meta"
            / "trex_track_force_prepare"
        ).glob("run_*/task_*.json"),
        key=lambda path: path.stat().st_mtime_ns,
    )
    entries: dict[int, dict] = {}
    for path in manifests:
        try:
            payload = _read_json(path)
        except Exception:
            continue
        for key, entry in payload.get("episodes", {}).items():
            if entry.get("status") == "complete":
                entries[int(key)] = copy.deepcopy(entry)
    expected = int(
        _read_json(source_force_root / "meta" / "info.json")["total_episodes"]
    )
    missing = [index for index in range(expected) if index not in entries]
    if missing:
        raise ValueError(
            f"source task manifests do not cover every episode: {missing[:10]}"
        )
    return entries


def _build_force_manifest(
    *,
    variant_root: Path,
    final_root: Path,
    registry: Sequence[EpisodeRecord],
    source_entries: dict[int, dict],
    track_hashes: dict[int, str],
) -> dict:
    final_track_cache = final_root / TRACK_CACHE_NAME
    manifest = force_builder._new_manifest(final_root, final_track_cache)
    manifest["filter_provenance"] = {
        "source_to_new_map": "meta/source_episode_index_map.json",
        "excluded_source_indices": "meta/excluded_source_episode_indices.json",
    }
    for record in registry:
        entry = copy.deepcopy(source_entries[record.source_episode_index])
        entry["episode_index"] = record.episode_index
        entry["source_episode_index"] = record.source_episode_index
        entry["num_frames"] = record.length
        entry["parquet"] = str(
            _episode_parquet_path(Path("."), record.episode_index)
        )
        entry["track_npz"] = str(
            final_track_cache / f"episode_{record.episode_index:06d}.npz"
        )
        entry["track_sha256"] = track_hashes[record.episode_index]
        entry["status"] = "complete"
        entry["skipped"] = False
        manifest["episodes"][f"{record.episode_index:06d}"] = entry
    _atomic_write_json(
        variant_root / "meta" / "trex_track_force_manifest.json",
        manifest,
    )
    return manifest


def _compute_base_stats(dataset_root: Path, registry: Sequence[EpisodeRecord]) -> dict:
    import pyarrow.parquet as pq

    stats: dict[str, dict] = {}
    for column in ("observation.state", "action", "timestamp"):
        parts: list[np.ndarray] = []
        for record in registry:
            table = pq.read_table(
                _episode_parquet_path(dataset_root, record.episode_index),
                columns=[column],
            )
            values = force_builder._column_to_numpy(table, column)
            if values.ndim == 1:
                values = values[:, None]
            parts.append(values)
        stats[column] = force_builder._statistics(np.concatenate(parts, axis=0))
    return stats


def _compute_base_relative_stats(
    dataset_root: Path,
    registry: Sequence[EpisodeRecord],
) -> dict:
    import pyarrow.parquet as pq

    action_offsets = range(24)
    output: dict[str, dict] = {}
    for name, selection in (
        ("left_arm", slice(0, 7)),
        ("right_arm", slice(29, 36)),
    ):
        parts: list[np.ndarray] = []
        for record in registry:
            table = pq.read_table(
                _episode_parquet_path(dataset_root, record.episode_index),
                columns=["observation.state", "action"],
            )
            state = force_builder._column_to_numpy(
                table,
                "observation.state",
            )[:, selection]
            action = force_builder._column_to_numpy(
                table,
                "action",
            )[:, selection]
            usable_length = len(state) - max(action_offsets)
            if usable_length <= 0:
                continue
            reference = state[:usable_length]
            parts.extend(
                action[offset : offset + usable_length] - reference
                for offset in action_offsets
            )
        if not parts:
            raise ValueError(f"no relative action samples for {name}")
        output[name] = force_builder._statistics(
            np.concatenate(parts, axis=0)
        )
    return output


def _finalize_metadata(
    *,
    variant_root: Path,
    registry: Sequence[EpisodeRecord],
    force: bool,
) -> None:
    stats_path = variant_root / "meta" / "stats.json"
    if force:
        base_variant_root = variant_root.parent / variant_root.name.removesuffix(
            "_force"
        )
        base_stats_path = base_variant_root / "meta" / "stats.json"
        if not base_stats_path.is_file():
            raise FileNotFoundError(
                f"base variant stats must be finalized first: {base_stats_path}"
            )
        _atomic_write_json(stats_path, _read_json(base_stats_path))
    else:
        _atomic_write_json(
            stats_path,
            _compute_base_stats(variant_root, registry),
        )
    if not force:
        _atomic_write_json(
            variant_root / "meta" / "relative_stats_dreamzero.json",
            _compute_base_relative_stats(variant_root, registry),
        )
        return
    result = force_builder.update_metadata(
        variant_root,
        assume_all_converted=True,
    )
    manifest_path = (
        variant_root / "meta" / "trex_track_force_manifest.json"
    )
    manifest = _read_json(manifest_path)
    manifest["metadata"] = result
    manifest["updated_at"] = _utc_now()
    _atomic_write_json(manifest_path, manifest)
    force_builder.validate_metadata(variant_root)
    for backup in (variant_root / "meta").glob("*.trex_track_force.bak"):
        backup.unlink()


def _copy_audit(source_force_root: Path, destination_force_root: Path) -> None:
    source = source_force_root / "audit"
    destination = destination_force_root / "audit" / "source_dataset"
    if destination.exists() or not source.exists():
        return
    shutil.copytree(source, destination)


def _validate_variant(
    *,
    root: Path,
    expected_registry: Sequence[EpisodeRecord],
    video_keys: Sequence[str],
    force: bool,
) -> dict:
    import pyarrow.parquet as pq

    info = _read_json(root / "meta" / "info.json")
    episodes = _read_jsonl(root / "meta" / "episodes.jsonl")
    tasks = _read_jsonl(root / "meta" / "tasks.jsonl")
    if int(info["total_episodes"]) != len(expected_registry):
        raise ValueError(f"{root}: total_episodes is stale")
    if int(info["total_frames"]) != sum(r.length for r in expected_registry):
        raise ValueError(f"{root}: total_frames is stale")
    if int(info["total_tasks"]) != len(tasks):
        raise ValueError(f"{root}: total_tasks is stale")
    if int(info["total_videos"]) != len(expected_registry) * len(video_keys):
        raise ValueError(f"{root}: total_videos is stale")
    if [int(row["episode_index"]) for row in episodes] != list(
        range(len(expected_registry))
    ):
        raise ValueError(f"{root}: episodes.jsonl is not dense")
    if [int(row["task_index"]) for row in tasks] != list(range(len(tasks))):
        raise ValueError(f"{root}: tasks.jsonl is not dense")

    manifest = None
    if force:
        force_builder.validate_metadata(root)
        manifest = _read_json(
            root / "meta" / "trex_track_force_manifest.json"
        )
        if len(manifest.get("episodes", {})) != len(expected_registry):
            raise ValueError(f"{root}: force manifest count is stale")

    for completed, record in enumerate(expected_registry, start=1):
        path = _episode_parquet_path(root, record.episode_index)
        if not _parquet_is_ready(
            path,
            episode_index=record.episode_index,
            task_index=record.task_index,
            global_index_start=record.global_index_start,
            expected_rows=record.length,
            force=force,
        ):
            raise ValueError(f"{root}: invalid parquet {path}")
        for video_key in video_keys:
            video_path = _episode_video_path(
                root,
                record.episode_index,
                video_key,
            )
            if not video_path.is_file() or video_path.stat().st_size <= 0:
                raise ValueError(f"{root}: missing video {video_path}")
        if force:
            track_path = _track_path(root, record.episode_index)
            if not _track_is_ready(
                track_path,
                episode_index=record.episode_index,
                expected_frames=record.length,
            ):
                raise ValueError(f"{root}: invalid track cache {track_path}")
            entry = manifest["episodes"].get(
                f"{record.episode_index:06d}",
                {},
            )
            if (
                entry.get("status") != "complete"
                or int(entry.get("source_episode_index", -1))
                != record.source_episode_index
            ):
                raise ValueError(
                    f"{root}: invalid manifest entry {record.episode_index}"
                )
        if completed % 250 == 0 or completed == len(expected_registry):
            print(
                f"validate {root.name}: {completed}/{len(expected_registry)}",
                flush=True,
            )

    sample_indices = sorted(
        {
            0,
            len(expected_registry) // 2,
            len(expected_registry) - 1,
        }
    )
    for episode_index in sample_indices:
        parquet_file = pq.ParquetFile(
            _episode_parquet_path(root, episode_index)
        )
        if int(parquet_file.metadata.num_rows) <= 0:
            raise ValueError(f"{root}: empty sample parquet {episode_index}")
        if force:
            force_builder.validate_episode_parquet(
                _episode_parquet_path(root, episode_index),
                verify_source_fk=False,
            )
    result = {
        "root": str(root),
        "episodes": len(expected_registry),
        "frames": int(sum(record.length for record in expected_registry)),
        "tasks": len(tasks),
        "videos": len(expected_registry) * len(video_keys),
        "force": force,
        "validated_at": _utc_now(),
    }
    _atomic_write_json(root / "meta" / "dataset_ready.json", result)
    return result


def _stage_paths(stage_root: Path) -> dict[str, Path]:
    return {name: stage_root / name for name in VARIANT_NAMES}


def _load_inputs(args: argparse.Namespace):
    source_info = _read_json(args.base_root / "meta" / "info.json")
    source_modality = _read_json(args.base_root / "meta" / "modality.json")
    source_embodiment = _read_json(args.base_root / "meta" / "embodiment.json")
    episodes = _read_jsonl(args.base_root / "meta" / "episodes.jsonl")
    source_tasks = _read_jsonl(args.base_root / "meta" / "tasks.jsonl")
    excluded = set(json.loads(args.blacklist.read_text()))
    if len(excluded) != 115:
        raise ValueError(
            f"expected 115 excluded episodes, got {len(excluded)}"
        )
    full_registry, full_tasks = _build_registry(
        episodes,
        source_tasks,
        excluded=excluded,
        source_limit=None,
    )
    small_registry, small_tasks = _build_registry(
        episodes,
        source_tasks,
        excluded=excluded,
        source_limit=100,
    )
    if len(full_registry) != 5349 or len(small_registry) != 100:
        raise ValueError(
            f"unexpected registry sizes: full={len(full_registry)}, "
            f"small={len(small_registry)}"
        )
    return (
        source_info,
        source_modality,
        source_embodiment,
        episodes,
        source_tasks,
        excluded,
        full_registry,
        full_tasks,
        small_registry,
        small_tasks,
    )


def _build(args: argparse.Namespace) -> None:
    (
        source_info,
        source_modality,
        source_embodiment,
        _,
        _,
        excluded,
        full_registry,
        full_tasks,
        small_registry,
        small_tasks,
    ) = _load_inputs(args)
    paths = _stage_paths(args.stage_root)
    video_keys = _video_keys(source_info)
    blacklist_sha256 = _sha256(args.blacklist)
    final_paths = {
        name: args.data_root / name
        for name in VARIANT_NAMES
    }
    args.stage_root.mkdir(parents=True, exist_ok=True)

    for name, registry, tasks, force, source_dataset in (
        (
            "trex_full",
            full_registry,
            full_tasks,
            False,
            args.base_root,
        ),
        (
            "trex_full_force",
            full_registry,
            full_tasks,
            True,
            args.force_root,
        ),
        (
            "trex_small",
            small_registry,
            small_tasks,
            False,
            args.base_root,
        ),
        (
            "trex_small_force",
            small_registry,
            small_tasks,
            True,
            args.force_root,
        ),
    ):
        _write_variant_metadata_skeleton(
            variant_root=paths[name],
            final_root=final_paths[name],
            source_info=source_info,
            source_modality=source_modality,
            source_embodiment=source_embodiment,
            registry=registry,
            tasks=tasks,
            video_keys=video_keys,
            excluded_source_indices=(
                sorted(excluded) if "full" in name else []
            ),
            force=force,
            source_dataset=source_dataset,
            blacklist_sha256=blacklist_sha256,
        )

    _rewrite_parquets(
        source_root=args.base_root,
        destination_root=paths["trex_full"],
        registry=full_registry,
        workers=args.parquet_workers,
        force=False,
    )
    _link_videos(
        source_root=args.base_root,
        destination_root=paths["trex_full"],
        registry=full_registry,
        video_keys=video_keys,
        source_uses_new_indices=False,
        workers=args.link_workers,
    )

    _rewrite_parquets(
        source_root=args.force_root,
        destination_root=paths["trex_full_force"],
        registry=full_registry,
        workers=args.parquet_workers,
        force=True,
    )
    _link_videos(
        source_root=paths["trex_full"],
        destination_root=paths["trex_full_force"],
        registry=full_registry,
        video_keys=video_keys,
        source_uses_new_indices=True,
        workers=args.link_workers,
    )
    full_track_hashes = _rewrite_tracks(
        source_root=args.force_root,
        destination_root=paths["trex_full_force"],
        registry=full_registry,
        workers=args.track_workers,
    )

    source_entries = _aggregate_source_manifest_entries(args.force_root)
    _build_force_manifest(
        variant_root=paths["trex_full_force"],
        final_root=final_paths["trex_full_force"],
        registry=full_registry,
        source_entries=source_entries,
        track_hashes=full_track_hashes,
    )
    _copy_audit(args.force_root, paths["trex_full_force"])

    for small_record, full_record in zip(
        small_registry,
        full_registry[: len(small_registry)],
    ):
        if (
            small_record.source_episode_index
            != full_record.source_episode_index
            or small_record.episode_index != full_record.episode_index
            or small_record.task_index != full_record.task_index
            or small_record.global_index_start
            != full_record.global_index_start
        ):
            raise ValueError("small is not an identity prefix of full")

    _link_subset_files(
        source_root=paths["trex_full"],
        destination_root=paths["trex_small"],
        registry=small_registry,
        video_keys=video_keys,
        force=False,
        workers=args.link_workers,
    )
    small_track_hashes = _link_subset_files(
        source_root=paths["trex_full_force"],
        destination_root=paths["trex_small_force"],
        registry=small_registry,
        video_keys=video_keys,
        force=True,
        workers=args.link_workers,
    )
    _build_force_manifest(
        variant_root=paths["trex_small_force"],
        final_root=final_paths["trex_small_force"],
        registry=small_registry,
        source_entries=source_entries,
        track_hashes=small_track_hashes,
    )
    _atomic_write_json(
        args.stage_root / "build_complete.json",
        {
            "completed_at": _utc_now(),
            "variants": list(VARIANT_NAMES),
        },
    )


def _metadata(args: argparse.Namespace) -> None:
    (
        _,
        _,
        _,
        _,
        _,
        _,
        full_registry,
        _,
        small_registry,
        _,
    ) = _load_inputs(args)
    paths = _stage_paths(args.stage_root)
    selected = set(args.metadata_variants)
    for name, registry, force in (
        ("trex_small", small_registry, False),
        ("trex_small_force", small_registry, True),
        ("trex_full", full_registry, False),
        ("trex_full_force", full_registry, True),
    ):
        if name not in selected:
            continue
        print(f"Finalizing metadata: {name}", flush=True)
        _finalize_metadata(
            variant_root=paths[name],
            registry=registry,
            force=force,
        )
    _atomic_write_json(
        args.stage_root / "metadata_complete.json",
        {"completed_at": _utc_now()},
    )


def _validate(args: argparse.Namespace) -> None:
    (
        source_info,
        _,
        _,
        _,
        _,
        _,
        full_registry,
        _,
        small_registry,
        _,
    ) = _load_inputs(args)
    paths = _stage_paths(args.stage_root)
    video_keys = _video_keys(source_info)
    results = []
    for name, registry, force in (
        ("trex_small", small_registry, False),
        ("trex_small_force", small_registry, True),
        ("trex_full", full_registry, False),
        ("trex_full_force", full_registry, True),
    ):
        results.append(
            _validate_variant(
                root=paths[name],
                expected_registry=registry,
                video_keys=video_keys,
                force=force,
            )
        )
    _atomic_write_json(
        args.stage_root / "validation_complete.json",
        {
            "validated_at": _utc_now(),
            "results": results,
        },
    )


def _install(args: argparse.Namespace) -> None:
    validation_marker = args.stage_root / "validation_complete.json"
    if not validation_marker.is_file():
        raise RuntimeError("staging validation marker is missing")
    paths = _stage_paths(args.stage_root)
    if any(not path.is_dir() for path in paths.values()):
        raise RuntimeError("one or more staged variants are missing")

    backup_root = args.data_root / (
        f".trex_old_{datetime.now().strftime('%Y%m%dT%H%M%S')}"
    )
    backup_root.mkdir(parents=True, exist_ok=False)
    moved_old: list[tuple[Path, Path]] = []
    installed: list[tuple[Path, Path]] = []
    try:
        for name in dict.fromkeys((*LEGACY_NAMES, *VARIANT_NAMES)):
            current = args.data_root / name
            if current.exists():
                backup = backup_root / name
                os.replace(current, backup)
                moved_old.append((backup, current))
        for name in VARIANT_NAMES:
            staged = paths[name]
            final = args.data_root / name
            os.replace(staged, final)
            installed.append((final, staged))

        (
            source_info,
            _,
            _,
            _,
            _,
            _,
            full_registry,
            _,
            small_registry,
            _,
        ) = _load_inputs_from_backups(args, backup_root)
        video_keys = _video_keys(source_info)
        for name, registry, force in (
            ("trex_small", small_registry, False),
            ("trex_small_force", small_registry, True),
            ("trex_full", full_registry, False),
            ("trex_full_force", full_registry, True),
        ):
            _validate_variant(
                root=args.data_root / name,
                expected_registry=registry,
                video_keys=video_keys,
                force=force,
            )
    except Exception:
        for final, staged in reversed(installed):
            if final.exists():
                os.replace(final, staged)
        for backup, current in reversed(moved_old):
            if backup.exists():
                os.replace(backup, current)
        backup_root.rmdir()
        raise

    shutil.rmtree(backup_root)
    args.stage_root.mkdir(parents=True, exist_ok=True)
    for marker in args.stage_root.glob("*.json"):
        marker.unlink()
    args.stage_root.rmdir()


def _load_inputs_from_backups(args: argparse.Namespace, backup_root: Path):
    backup_args = copy.copy(args)
    backup_args.base_root = backup_root / "trex_datasetv2"
    backup_args.force_root = backup_root / "trex_track_force_v2"
    backup_args.blacklist = (
        backup_args.force_root
        / "audit"
        / "track_quality"
        / "frozen_wrist_episode_indices.json"
    )
    return _load_inputs(backup_args)


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--data-root",
        type=Path,
        default=_DREAMZERO_ROOT / "data",
    )
    parser.add_argument(
        "--base-root",
        type=Path,
        default=_DREAMZERO_ROOT / "data" / "trex_datasetv2",
    )
    parser.add_argument(
        "--force-root",
        type=Path,
        default=_DREAMZERO_ROOT / "data" / "trex_track_force_v2",
    )
    parser.add_argument(
        "--blacklist",
        type=Path,
        default=(
            _DREAMZERO_ROOT
            / "data"
            / "trex_track_force_v2"
            / "audit"
            / "track_quality"
            / "frozen_wrist_episode_indices.json"
        ),
    )
    parser.add_argument(
        "--stage-root",
        type=Path,
        default=_DREAMZERO_ROOT / "data" / ".trex_variants_staging",
    )
    parser.add_argument(
        "--phase",
        choices=("build", "metadata", "validate", "install", "all"),
        default="all",
    )
    parser.add_argument("--parquet-workers", type=int, default=8)
    parser.add_argument("--track-workers", type=int, default=8)
    parser.add_argument("--link-workers", type=int, default=32)
    parser.add_argument(
        "--metadata-variants",
        nargs="+",
        choices=VARIANT_NAMES,
        default=list(VARIANT_NAMES),
        help="variants to process during the metadata phase",
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = _build_parser().parse_args(argv)
    for name in ("data_root", "base_root", "force_root", "blacklist", "stage_root"):
        setattr(args, name, getattr(args, name).expanduser().resolve())
    if args.phase in ("build", "all"):
        _build(args)
    if args.phase in ("metadata", "all"):
        _metadata(args)
    if args.phase in ("validate", "all"):
        _validate(args)
    if args.phase in ("install", "all"):
        _install(args)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())