#!/usr/bin/env python3
"""merge_raw_episodes.py — いくつかの raw_episodes を 1 つのデータセットにまとめる (2026-10-31 依頼 45).

お手本 (teach_v1_trim2) と補正エピソード (corr_v1) を合わせて, 1 つの raw_episodes にする.
1 つにするのは, 学習 (convert_to_lerobot.py --raw-dir) も実機 (run_policy_client.py --raw-dir の安全範囲・
eval_policy_offline.py) も「学習に使ったデータ」を 1 つのフォルダで受け取るため.

  python3 pi05/merge_raw_episodes.py \\
      --src mecharm_redblock_box_teach_v1_trim2 --src mecharm_redblock_box_corr_v1 \\
      --dataset mecharm_redblock_box_teach_v1_trim2_corr_v1 --tar

出力:
  datasets/{DATASET}/raw_episodes/{エピソード}/  元のエピソードの写し (frames/ episode.npz meta.json)
  datasets/{DATASET}/index.json / README.md     どの元から何本・何コマ
  datasets/{DATASET}/code/merge_raw_episodes.py このスクリプトの写し (データと対のコード)
  datasets/{DATASET}/{DATASET}_raw_episodes.tar (--tar) Pod へ送る形. ._* (AppleDouble) は入れない

決まり:
  - 元のデータセットは読むだけ (上書き・移動・削除しない)
  - 写す前に, 全部の元が お手本 (最初の --src) と一緒に使えるか (キー・型・FPS・画像の大きさ) を確かめ,
    合わなければ何も書かずに止まる
  - エピソード名が重なったら止まる. 出力先がもうあれば止まる (--force で作り直す)
  - exFAT に macOS が作る ._* (AppleDouble) は写さない
"""
import argparse
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

_HERE = Path(__file__).resolve()
sys.path.insert(0, str(_HERE.parent))

import build_demo_episodes as bde  # noqa: E402
from build_correction_episodes import _rmtree, check_compat  # noqa: E402

DATA_ROOT = bde.DATA_ROOT


def episodes(raw):
    return sorted(p.parent for p in Path(raw).glob("[!.]*/episode.npz"))


def ep_summary(ep):
    meta = json.loads((ep / "meta.json").read_text())
    return {"steps": int(meta["steps"]), "correct_steps": int(meta.get("correct_steps", 0)),
            "model_steps": int(meta.get("model_steps", 0)), "source": meta.get("source", "")}


def _ignore_appledouble(_dir, names):
    return [n for n in names if n.startswith("._")]


def make_tar(ds, name):
    """Pod へ送る tar (中身は raw_episodes/ だけ). ._* と macOS の拡張属性を入れない."""
    out = ds / f"{name}_raw_episodes.tar"
    env = dict(os.environ, COPYFILE_DISABLE="1")
    subprocess.run(["tar", "--exclude", "._*", "-cf", str(out), "raw_episodes"],
                   cwd=str(ds), env=env, check=True)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", action="append", required=True,
                    help="まとめる元のデータセット名 (datasets/ の下). 最初の 1 つをお手本として形をそろえる")
    ap.add_argument("--dataset", required=True, help="まとめた先のデータセット名 (新しい名前)")
    ap.add_argument("--force", action="store_true", help="まとめた先を作り直す (元には触らない)")
    ap.add_argument("--tar", action="store_true", help="Pod へ送る tar も作る")
    ap.add_argument("--data-root", help=argparse.SUPPRESS)   # テスト用 (KIOXIA の代わり)
    args = ap.parse_args()
    if args.data_root:
        root = Path(args.data_root)
    else:
        bde.ensure_ssd()
        root = DATA_ROOT
    dsroot = root / "datasets"
    if args.dataset in args.src:
        sys.exit(f"[ABORT] まとめた先が元と同じ名前: {args.dataset}")
    srcs = [dsroot / s / "raw_episodes" for s in args.src]
    for s in srcs:
        if not episodes(s):
            sys.exit(f"[ABORT] エピソードがない: {s}")

    # 形の確かめ (書く前に)
    ref = srcs[0]
    for s in srcs[1:]:
        probs = check_compat(s, ref)
        if probs:
            print(f"[ABORT] {s} がお手本 {ref} と合わない:")
            for p in probs:
                print("        " + p)
            return 1
    names = {}
    for s in srcs:
        for ep in episodes(s):
            if ep.name in names:
                sys.exit(f"[ABORT] エピソード名が重なる: {ep.name} ({names[ep.name]} と {s})")
            names[ep.name] = s

    ds = dsroot / args.dataset
    if ds.exists():
        if not args.force:
            sys.exit(f"[ABORT] もうある: {ds} (作り直すなら --force)")
        _rmtree(ds)
    out = ds / "raw_episodes"
    out.mkdir(parents=True)
    index = {"sources": {}, "episodes": {}}
    for src_name, s in zip(args.src, srcs):
        eps = episodes(s)
        for ep in eps:
            shutil.copytree(ep, out / ep.name, ignore=_ignore_appledouble)
            index["episodes"][ep.name] = dict(ep_summary(ep), from_dataset=src_name)
        index["sources"][src_name] = {
            "episodes": len(eps), "steps": sum(index["episodes"][e.name]["steps"] for e in eps),
            "correct_steps": sum(index["episodes"][e.name]["correct_steps"] for e in eps)}
        print(f"[COPY] {src_name}: {len(eps)} 本 / {index['sources'][src_name]['steps']} コマ")

    # 写したあとも同じ形か (コマ数・キー)
    probs = check_compat(out, ref)
    if probs:
        print("[NG  ] まとめたあとの確かめ:")
        for p in probs:
            print("        " + p)
        return 1
    total_eps = len(index["episodes"])
    total_steps = sum(v["steps"] for v in index["episodes"].values())
    index.update(total_episodes=total_eps, total_steps=total_steps,
                 repo_id=f"local/{args.dataset}")
    (ds / "index.json").write_text(json.dumps(index, indent=2, ensure_ascii=False))
    lines = [
        f"# {args.dataset}", "",
        "いくつかの raw_episodes を 1 つにまとめたデータセット (元は読むだけ。写しを置く)。", "",
        "- 作成: `python3 pi05/merge_raw_episodes.py "
        + " ".join(f"--src {s}" for s in args.src) + f" --dataset {args.dataset}`",
        f"- LeRobot に変換するときの repo_id: `local/{args.dataset}`", "",
        f"合計 {total_eps} 本 / {total_steps} コマ", "",
        "| 元 | 本数 | コマ | うち補正 (correct) |", "| --- | --- | --- | --- |",
    ]
    for k, v in index["sources"].items():
        lines.append(f"| {k} | {v['episodes']} | {v['steps']} | {v['correct_steps']} |")
    (ds / "README.md").write_text("\n".join(lines) + "\n")
    code = ds / "code"
    code.mkdir(exist_ok=True)
    shutil.copy2(_HERE, code / _HERE.name)   # データと対のコード
    print(f"\n合計 {total_eps} 本 / {total_steps} コマ → {out}")
    if args.tar:
        t = make_tar(ds, args.dataset)
        print(f"[TAR ] {t} ({t.stat().st_size / 1e6:.0f} MB)")
    return 0


if __name__ == "__main__":
    sys.exit(main())