#!/usr/bin/env python3
"""Extract a low-variation Bridge/LeRobot subset by episode task text.

The default selection is the plate-placement family (``put ... on plate``),
which is close to the current carrot-on-plate evaluation while retaining
multiple objects and camera scenes.  Episodes are copied whole, so no video
re-encoding or temporal duplication is introduced.
"""
from __future__ import annotations

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


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default="/opt/zhangchenyu/datasets/bridge_orig_lerobot_smoke")
    ap.add_argument("--dst", default=".artifacts/bridge_put_on_plate_259")
    ap.add_argument("--pattern", default=r"^put .+ on plate(?: sink)?$")
    ap.add_argument("--max-episodes", type=int, default=None)
    ap.add_argument("--min-length", type=int, default=None,
                     help="keep only episodes with at least this many frames")
    ap.add_argument("--hardlink-videos", action="store_true",
                     help="hardlink videos instead of duplicating them when possible")
    args = ap.parse_args()
    src, dst = Path(args.src), Path(args.dst)
    meta = src / "meta"
    episodes = [json.loads(x) for x in (meta / "episodes.jsonl").read_text().splitlines()]
    tasks = {int(x["task_index"]): x["task"] for x in
             (json.loads(x) for x in (meta / "tasks.jsonl").read_text().splitlines())}
    rx = re.compile(args.pattern, re.I)
    selected = []
    for ep in episodes:
        text = " ".join(ep.get("tasks", []))
        if rx.fullmatch(text.strip()):
            selected.append((ep, text.strip()))
    selected.sort(key=lambda x: int(x[0]["episode_index"]))
    if args.min_length is not None:
        selected = [x for x in selected if int(x[0]["length"]) >= args.min_length]
    if args.max_episodes is not None:
        selected = selected[:args.max_episodes]
    if not selected:
        raise SystemExit("no matching episodes")

    missing = []
    for ep, _ in selected:
        i = int(ep["episode_index"])
        chunk = i // 1000
        dp = src / f"data/chunk-{chunk:03d}/episode_{i:06d}.parquet"
        vp = src / f"videos/chunk-{chunk:03d}/observation.images.image_0/episode_{i:06d}.mp4"
        if not dp.exists() or not vp.exists():
            missing.extend(str(p) for p in (dp, vp) if not p.exists())
    if missing:
        raise SystemExit("raw parquet/video files are missing; metadata-only source. First missing: " + missing[0])

    # Import only after the source check: this script remains usable for dry inspection.
    import pyarrow.parquet as pq
    import pyarrow as pa

    (dst / "meta").mkdir(parents=True, exist_ok=True)
    out_eps, out_tasks = [], []
    total_frames = 0
    for new_i, (ep, text) in enumerate(selected):
        old_i = int(ep["episode_index"])
        chunk = old_i // 1000
        old_dp = src / f"data/chunk-{chunk:03d}/episode_{old_i:06d}.parquet"
        old_vp = src / f"videos/chunk-{chunk:03d}/observation.images.image_0/episode_{old_i:06d}.mp4"
        new_dp = dst / f"data/chunk-000/episode_{new_i:06d}.parquet"
        new_vp = dst / f"videos/chunk-000/observation.images.image_0/episode_{new_i:06d}.mp4"
        new_dp.parent.mkdir(parents=True, exist_ok=True)
        new_vp.parent.mkdir(parents=True, exist_ok=True)
        tab = pq.read_table(old_dp)
        # Keep all rows but rewrite episode/task ids to a compact standalone set.
        cols = {name: tab[name] for name in tab.column_names}
        if "episode_index" in cols:
            cols["episode_index"] = pa.array([new_i] * tab.num_rows, type=cols["episode_index"].type)
        task_col = pa.array([new_i] * tab.num_rows, type=cols["task_index"].type)
        cols["task_index"] = task_col
        pq.write_table(pa.table(cols), new_dp)
        if args.hardlink_videos:
            os.link(old_vp, new_vp)
        else:
            shutil.copy2(old_vp, new_vp)
        out_eps.append({"episode_index": new_i, "tasks": [text], "length": int(ep["length"])})
        out_tasks.append({"task_index": new_i, "task": text})
        total_frames += int(ep["length"])

    info = json.loads((meta / "info.json").read_text())
    info["total_episodes"] = len(out_eps)
    info["total_frames"] = total_frames
    info["total_tasks"] = len(out_tasks)
    info["total_videos"] = len(out_eps)
    info["total_chunks"] = 1
    info["splits"] = {"train": f"0:{len(out_eps)}"}
    info["data_path"] = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet"
    info["video_path"] = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4"
    (dst / "meta/info.json").write_text(json.dumps(info, indent=2) + "\n")
    (dst / "meta/episodes.jsonl").write_text("".join(json.dumps(x) + "\n" for x in out_eps))
    (dst / "meta/tasks.jsonl").write_text("".join(json.dumps(x) + "\n" for x in out_tasks))
    for name in ("modality.json", "stats.json", "stats_gr00t.json"):
        if (meta / name).exists(): shutil.copy2(meta / name, dst / "meta" / name)
    print(json.dumps({"episodes": len(out_eps), "frames": total_frames, "dst": str(dst)}, indent=2))


if __name__ == "__main__":
    main()