#!/usr/bin/env python
"""One-time shard builder: LeRobot v3 video dataset -> sequential JPEG shards.

Why this exists
---------------
Training reads ONE random frame per sample. On a LeRobot v3 repo that means an
h264 seek into a ~195 MB mp4 holding ~10 concatenated episodes, per camera, per
sample. That is the bottleneck the project notes already call out ("video decode
is the training bottleneck, use 12+ dataloader workers"), and no amount of GPU
fixes it.

This pass decodes every video ONCE, sequentially (the cheap direction), at a
target rate, resizes to the training resolution, and writes JPEG bytes into
parquet shards. Training then does sequential reads + JPEG decode, which is
~100x cheaper per frame and trivially parallel.

Layout produced in --out:
    shards/shard-{i:05d}.parquet   episode_index, frame_index, cam0, cam1 (JPEG bytes)
    frames.parquet                 episode_index, frame_index, state[D], action[D], task_index
    tasks.json                     task_index -> instruction string
    stats.json                     per-source mean/std for state and action
    manifest.json                  fps, image_size, dims, camera keys, counts

Usage (parallel across boxes: each takes a disjoint episode range):
    python build_shards.py --repo-id zekaiwang/trex_dataset --out /data/trex_shards \
        --hz 10 --image-size 256 --part 0 --num-parts 4

Verified against zekaiwang/trex_dataset metadata: episodes are contiguous and
ordered inside every mp4, frame counts match the timestamps exactly, and the
declared fps (30) is the real one.
"""

from __future__ import annotations

import argparse
import io
import json
import os
from pathlib import Path

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq


# --------------------------------------------------------------------- metadata


def load_metadata(repo_id: str, revision: str | None):
    from huggingface_hub import hf_hub_download

    def grab(path):
        return hf_hub_download(repo_id, path, repo_type="dataset", revision=revision)

    info = json.loads(Path(grab("meta/info.json")).read_text())

    from huggingface_hub import HfApi

    api = HfApi()
    files = api.list_repo_files(repo_id, repo_type="dataset", revision=revision)
    ep_files = sorted(f for f in files if f.startswith("meta/episodes/") and f.endswith(".parquet"))
    tables = [pq.read_table(grab(f)) for f in ep_files]
    episodes = pa.concat_tables(tables) if len(tables) > 1 else tables[0]
    return info, episodes, files


def pick_cameras(info: dict, cam0: str | None, cam1: str | None):
    """Mirror CanonicalSource's cam0/cam1 choice unless the caller overrides it.

    CanonicalSource sorts the image keys, takes PRIMARY_RE (or the first key) as
    cam0 and the first WRIST_RE match as cam1. On trex that resolves to
    head_left / left_wrist — right_wrist and all 20 tactile streams are dropped.
    Made explicit here rather than left to a regex.
    """
    import re

    PRIMARY_RE = re.compile(r"top|front|base|main|scene|image_0$|images\.image$|image2$")
    WRIST_RE = re.compile(r"wrist|gripper|hand|image_1$")
    keys = sorted(k for k, v in info["features"].items() if v.get("dtype") == "video")
    if cam0 is None:
        cam0 = next((k for k in keys if PRIMARY_RE.search(k)), keys[0])
    if cam1 is None:
        cam1 = next((k for k in keys if k != cam0 and WRIST_RE.search(k)), None)
    assert cam0 in keys, f"{cam0} not among {keys}"
    assert cam1 is None or cam1 in keys, f"{cam1} not among {keys}"
    return cam0, cam1


# ------------------------------------------------------------------ video decode


class VideoReader:
    """Opens one mp4 at a time, seeks to an episode, decodes its frames forward.

    Episodes inside a file are contiguous and time-ordered (verified), so walking
    episodes in order keeps this effectively sequential: one seek per episode,
    costing at most one extra GOP.
    """

    def __init__(self, repo_id, revision, key, cache_dir, evict: bool):
        self.repo_id, self.revision, self.key = repo_id, revision, key
        self.cache_dir, self.evict = cache_dir, evict
        self._path = None
        self._ident = None
        self._container = None
        self._stream = None

    def _open(self, chunk_index: int, file_index: int):
        ident = (chunk_index, file_index)
        if ident == self._ident:
            return
        self._close(drop=True)
        from huggingface_hub import hf_hub_download

        rel = f"videos/{self.key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4"
        self._path = hf_hub_download(
            self.repo_id, rel, repo_type="dataset", revision=self.revision, cache_dir=self.cache_dir
        )
        import av

        self._container = av.open(self._path)
        self._stream = self._container.streams.video[0]
        self._stream.thread_type = "AUTO"  # frame+slice threading; the whole point
        self._ident = ident

    def _close(self, drop: bool):
        if self._container is not None:
            self._container.close()
            self._container = None
        if drop and self.evict and self._path and os.path.exists(self._path):
            real = os.path.realpath(self._path)  # hf cache stores a symlink to the blob
            for p in {self._path, real}:
                try:
                    os.remove(p)
                except OSError:
                    pass
        self._path = self._ident = self._stream = None

    def close(self):
        self._close(drop=True)

    def read_episode(self, chunk_index, file_index, from_ts, n_frames, fps, stride):
        """Yield (frame_index_in_episode, PIL.Image) for every stride-th frame."""
        self._open(chunk_index, file_index)
        tb = self._stream.time_base
        target = int(from_ts / tb)
        self._container.seek(target, stream=self._stream, backward=True, any_frame=False)
        wanted = set(range(0, n_frames, stride))
        eps = 0.5 / fps  # half a frame of slack against float timestamps
        got = 0
        for frame in self._container.decode(self._stream):
            if frame.pts is None:
                continue
            t = float(frame.pts * tb)
            if t < from_ts - eps:
                continue
            idx = int(round((t - from_ts) * fps))
            if idx >= n_frames:
                break
            if idx in wanted:
                wanted.discard(idx)
                got += 1
                yield idx, frame.to_image()
            if not wanted:
                break
        if wanted:
            raise RuntimeError(
                f"{self.key} chunk{chunk_index}/file{file_index} @{from_ts:.3f}: "
                f"decoded {got}, still missing {len(wanted)} frames"
            )


def encode_jpeg(img, size: int, quality: int) -> bytes:
    """Resize to (size, size) and JPEG-encode.

    Plain resize, no crop: this reproduces CanonicalSource._img, which squashes
    the aspect ratio with F.interpolate. Matching it matters more than being
    geometrically nicer, otherwise the shards are off-distribution for the
    pretrained checkpoint.
    """
    from PIL import Image

    if img.size != (size, size):
        img = img.resize((size, size), Image.BILINEAR)
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=quality, subsampling=0)
    return buf.getvalue()


# ------------------------------------------------------------------------ state


def load_frame_table(repo_id, revision, files, episodes, keep_mask, stride, cache_dir,
                     skip_bad=False):
    bad_rows: set[int] = set()
    """Pull state/action/task for the kept frames out of the data parquets."""
    from huggingface_hub import hf_hub_download

    ep = {c: np.asarray(episodes[c]) for c in
          ["episode_index", "length", "dataset_from_index", "data/chunk_index", "data/file_index"]}
    want = {}  # (chunk, file) -> list of episode rows
    for i in np.nonzero(keep_mask)[0]:
        want.setdefault((int(ep["data/chunk_index"][i]), int(ep["data/file_index"][i])), []).append(i)

    out_ep, out_fr, out_state, out_action, out_task = [], [], [], [], []
    tasks: dict[str, int] = {}
    for (ci, fi), rows in sorted(want.items()):
        rel = f"data/chunk-{ci:03d}/file-{fi:03d}.parquet"
        if rel not in files:
            raise FileNotFoundError(rel)
        path = hf_hub_download(repo_id, rel, repo_type="dataset", revision=revision, cache_dir=cache_dir)
        t = pq.read_table(path, columns=["episode_index", "frame_index", "observation.state", "action", "task_index"])
        d = {c: t[c].to_numpy(zero_copy_only=False) for c in t.column_names}
        for r in rows:
            e = int(ep["episode_index"][r])
            sel = np.nonzero(d["episode_index"] == e)[0]
            if len(sel) != int(ep["length"][r]):
                if skip_bad:
                    # unitree v3 repos ship a few episodes whose rows are
                    # missing from the indexed file — drop them consistently
                    # from BOTH the frame table and the video stage
                    print(f"  SKIP episode {e}: parquet has {len(sel)} rows, meta says {ep['length'][r]}")
                    bad_rows.add(int(r))
                    continue
                raise RuntimeError(f"episode {e}: parquet has {len(sel)} rows, meta says {ep['length'][r]}")
            order = np.argsort(d["frame_index"][sel])
            sel = sel[order][::stride]
            out_ep.append(np.full(len(sel), e, dtype=np.int32))
            out_fr.append(np.arange(len(sel), dtype=np.int32))
            out_state.append(np.stack([np.asarray(x, dtype=np.float32) for x in d["observation.state"][sel]]))
            out_action.append(np.stack([np.asarray(x, dtype=np.float32) for x in d["action"][sel]]))
            task = episodes["tasks"][r].as_py()
            task = task[0] if isinstance(task, list) and task else ""
            out_task.append(np.full(len(sel), tasks.setdefault(task, len(tasks)), dtype=np.int32))
        del t, d
    return (
        np.concatenate(out_ep), np.concatenate(out_fr),
        np.concatenate(out_state), np.concatenate(out_action),
        np.concatenate(out_task), tasks, bad_rows,
    )


# ------------------------------------------------------------------------- main


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo-id", default="zekaiwang/trex_dataset")
    ap.add_argument("--revision", default=None)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--hz", type=float, default=10.0, help="target rate; 0 = keep native fps")
    ap.add_argument("--image-size", type=int, default=256)
    ap.add_argument("--jpeg-quality", type=int, default=92)
    ap.add_argument("--rows-per-shard", type=int, default=20000)
    ap.add_argument("--cam0", default=None, help="override the primary camera key")
    ap.add_argument("--cam1", default=None, help="override the wrist camera key")
    ap.add_argument("--part", type=int, default=0)
    ap.add_argument("--num-parts", type=int, default=1)
    ap.add_argument("--max-episodes", type=int, default=None, help="smoke test on the first N")
    ap.add_argument("--cache-dir", default=None)
    ap.add_argument("--keep-videos", action="store_true", help="do NOT delete mp4s after use")
    ap.add_argument("--skip-bad-episodes", action="store_true",
                    help="drop episodes whose parquet rows disagree with meta instead of failing")
    args = ap.parse_args()

    info, episodes, files = load_metadata(args.repo_id, args.revision)
    fps = float(info["fps"])
    stride = 1 if args.hz <= 0 else max(1, int(round(fps / args.hz)))
    eff_hz = fps / stride
    cam0, cam1 = pick_cameras(info, args.cam0, args.cam1)

    n_ep = episodes.num_rows
    lengths = np.asarray(episodes["length"])
    # sanity: the LIBERO lesson — never trust declared fps
    for key in (cam0, cam1):
        if key is None:
            continue
        f = np.asarray(episodes[f"videos/{key}/from_timestamp"])
        t = np.asarray(episodes[f"videos/{key}/to_timestamp"])
        measured = lengths / (t - f)
        if not np.allclose(measured, fps, atol=1e-3):
            raise SystemExit(
                f"declared fps={fps} but {key} measures {measured.min():.3f}..{measured.max():.3f} — "
                "fix the rate before building shards"
            )

    order = np.arange(n_ep) if args.max_episodes is None else np.arange(min(n_ep, args.max_episodes))
    mine = order[args.part :: args.num_parts]
    keep = np.zeros(n_ep, dtype=bool)
    keep[mine] = True

    print(f"repo {args.repo_id}: {n_ep} episodes, {lengths.sum():,} frames @ {fps} Hz")
    print(f"cam0={cam0}\ncam1={cam1}")
    print(f"stride {stride} -> {eff_hz:g} Hz, part {args.part}/{args.num_parts}: {len(mine)} episodes")

    args.out.mkdir(parents=True, exist_ok=True)
    (args.out / "shards").mkdir(exist_ok=True)

    print("collecting state/action ...")
    ep_i, fr_i, state, action, task_i, tasks, bad_rows = load_frame_table(
        args.repo_id, args.revision, files, episodes, keep, stride, args.cache_dir,
        skip_bad=args.skip_bad_episodes,
    )
    if bad_rows:
        mine = np.asarray([r for r in mine if int(r) not in bad_rows])
        print(f"  dropped {len(bad_rows)} bad episodes, {len(mine)} remain")
    print(f"  {len(ep_i):,} frames kept, state {state.shape[1]}d, action {action.shape[1]}d, {len(tasks)} tasks")

    readers = {
        "cam0": VideoReader(args.repo_id, args.revision, cam0, args.cache_dir, not args.keep_videos),
        "cam1": None if cam1 is None else VideoReader(args.repo_id, args.revision, cam1, args.cache_dir, not args.keep_videos),
    }

    schema = pa.schema([
        ("episode_index", pa.int32()), ("frame_index", pa.int32()),
        ("cam0", pa.binary()), ("cam1", pa.binary()),
    ])
    buf = {k: [] for k in ["episode_index", "frame_index", "cam0", "cam1"]}
    shard_id, written, in_shard = args.part * 100000, 0, 0

    writer = None

    def flush(final=False):
        nonlocal writer, shard_id, buf, written, in_shard
        if buf["episode_index"]:
            if writer is None:
                path = args.out / "shards" / f"shard-{shard_id:05d}.parquet"
                writer = pq.ParquetWriter(path, schema, compression="zstd", compression_level=1)
            writer.write_table(pa.table(buf, schema=schema))
            written += len(buf["episode_index"])
            in_shard += len(buf["episode_index"])
            buf = {k: [] for k in buf}
        if writer is not None and (final or in_shard >= args.rows_per_shard):
            writer.close()
            writer = None
            shard_id += 1
            in_shard = 0

    import time

    t0 = time.time()
    for n, r in enumerate(mine):
        e = int(episodes["episode_index"][r].as_py())
        L = int(lengths[r])
        frames = {}
        for slot, key in (("cam0", cam0), ("cam1", cam1)):
            if key is None:
                continue
            rd = readers[slot]
            ci = int(episodes[f"videos/{key}/chunk_index"][r].as_py())
            fi = int(episodes[f"videos/{key}/file_index"][r].as_py())
            ts = float(episodes[f"videos/{key}/from_timestamp"][r].as_py())
            frames[slot] = {
                i // stride: encode_jpeg(img, args.image_size, args.jpeg_quality)
                for i, img in rd.read_episode(ci, fi, ts, L, fps, stride)
            }
        n_kept = len(range(0, L, stride))
        blank = b""
        for j in range(n_kept):
            buf["episode_index"].append(e)
            buf["frame_index"].append(j)
            buf["cam0"].append(frames["cam0"][j])
            buf["cam1"].append(frames.get("cam1", {}).get(j, blank))
        if len(buf["episode_index"]) >= 2048:
            flush()
        if (n + 1) % 50 == 0:
            el = time.time() - t0
            print(f"  {n+1}/{len(mine)} episodes, {written:,} frames, "
                  f"{written/el:.0f} fr/s, eta {(len(mine)-n-1)*el/(n+1)/60:.1f} min", flush=True)
    flush(final=True)
    for rd in readers.values():
        if rd is not None:
            rd.close()

    # frames.parquet holds only the kept rows, in shard order
    pq.write_table(
        pa.table({
            "episode_index": pa.array(ep_i, pa.int32()),
            "frame_index": pa.array(fr_i, pa.int32()),
            "state": pa.array(list(state), pa.list_(pa.float32(), state.shape[1])),
            "action": pa.array(list(action), pa.list_(pa.float32(), action.shape[1])),
            "task_index": pa.array(task_i, pa.int32()),
        }),
        args.out / f"frames-part{args.part:02d}.parquet",
        compression="zstd",
    )
    (args.out / f"tasks-part{args.part:02d}.json").write_text(
        json.dumps({v: k for k, v in tasks.items()}, ensure_ascii=False)
    )
    (args.out / f"stats-part{args.part:02d}.json").write_text(json.dumps({
        "observation.state": {"mean": state.mean(0).tolist(), "std": state.std(0).tolist()},
        "action": {"mean": action.mean(0).tolist(), "std": action.std(0).tolist()},
        "count": int(len(ep_i)),
    }))
    (args.out / f"manifest-part{args.part:02d}.json").write_text(json.dumps({
        "repo_id": args.repo_id, "revision": args.revision,
        "source_fps": fps, "stride": stride, "fps": eff_hz,
        "image_size": args.image_size, "jpeg_quality": args.jpeg_quality,
        "cam0": cam0, "cam1": cam1,
        "state_dim": int(state.shape[1]), "action_dim": int(action.shape[1]),
        "episodes": int(len(mine)), "frames": int(written),
    }, indent=2))
    print(f"done: {written:,} frames in {(time.time()-t0)/60:.1f} min -> {args.out}")


if __name__ == "__main__":
    main()