#!/usr/bin/env python3
"""build_correction_episodes.py — 実機推論 (巻き戻し・補正つき) の記録から π0.5 学習用エピソードを作る (2026-10-29).

入力 (1 試行 = 1 run_stem. run_policy_client.py --execute が残したもの):
  runs/pi05_policy/{run_stem}/{run_stem}.json            結果 (outcome / judgments)
  runs/pi05_policy/{run_stem}/log.jsonl                  推論・巻き戻し・補正の記録
  runs/pi05_policy/{run_stem}/recording/states.csv       actual (実測6軸) / gripper_cmd / stage (= 状態)
  runs/pi05_policy/{run_stem}/recording/top.mp4 + top_frame_log.csv

出力 (お手本 mecharm_redblock_box_teach_v1_trim2 と同じ形):
  datasets/{DATASET}/raw_episodes/{run_stem}/frames/000000.jpg ... / episode.npz / meta.json
  datasets/{DATASET}/index.json / README.md
  datasets/{DATASET}/code/  このスクリプトと読み込む部品の写し (データと対のコード)

作り方:
  - 使うのは 結果が「成功」(y) の run だけ
  - 時間軸は 10Hz. state は実測角の線形補間 + グリッパーは指令の時刻で 0/1 (build_demo_episodes と同じ)
    実測角の抜け (送信中) は --gap-fill で埋める (linear = corr_v1 / speed20 = corr_v2. 依頼 46)
  - 学習に使うのは状態が model / correct のステップだけ. rewind / home / wait (確認の入力待ち.
    人がブロックを置き直すと手が写る) / regrip (補正の R で開いてから閉じ終わるまで. 2026-10-30 依頼 35)
    のステップと,
    巻き戻しで取り消した区間 (区間を始めた時刻 〜 巻き戻しが終わった時刻) と,
    取り消した補正 (補正を始めた時刻 〜 戻り終えた時刻) は除く.
    除いたあとは「その手前」と「巻き戻したあとの続き」をつなぐ (戻った先は区間の始めの姿勢なので続く)
  - actions は「0.1 秒後の state」. 除いたあとに残ったステップで作り直す (trim2 と同じ作り方)
  - 止まっている区間は build_demo_episodes と同じ決まりで詰める (IDLE_* と --trim-mid-idle の
    mid_idle_keep_mask. 既定で詰める. --no-trim-mid-idle で前者だけ)
  - --trim-hold (corr_v3. 依頼 48): 中継点・補正の 1 刻みごとの「着いたあとの止まり」を, 区間ごとに
    残す数を変えて詰める (hold_keep_mask. 止まりの判定とずれの切り方は mid_idle_keep_mask と同じ)
  - 元の記録は読むだけ (上書き・移動・削除しない)

Usage:
  python3 pi05/build_correction_episodes.py                       # runs/pi05_policy の全 run
  python3 pi05/build_correction_episodes.py --run pi05_20261030_101500 --force
  python3 pi05/build_correction_episodes.py --check-with \\
      /Volumes/KIOXIA/r2r2r-data/datasets/mecharm_redblock_box_teach_v1_trim2/raw_episodes
"""
import argparse
import csv
import json
import shutil
import sys
from pathlib import Path

import cv2
import numpy as np

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

import build_demo_episodes as bde  # noqa: E402
from rewind_correct import MODE_CORRECT, MODE_MODEL, excluded_windows, step_keep_mask  # noqa: E402

DATA_ROOT = bde.DATA_ROOT
RUNS_ROOT = DATA_ROOT / "runs" / "pi05_policy"
DEFAULT_DATASET = "mecharm_redblock_box_corr_v1"
FPS = bde.FPS
# ---- 実測角の抜けの埋め方 (--gap-fill. 2026-10-31 依頼 46) ----
# 記録スレッドと本体は daemon への 1 本のつながりを鍵で順に使う. send_angles は daemon が送ったあとに
# 角度を読み直す間 (約 1.6 秒. 動いている間は読めずにやり直す) 返ってこないので, その間は実測角が残らない.
# 腕はこの抜けの頭で動き出し, 抜けの終わりの実測角 (着いた姿勢) で止まっている.
#   linear  : 抜けの前後の実測角をまっすぐ結ぶ (corr_v1. 1.6 秒かけてゆっくり動いたことになる)
#   speed20 : 抜けの頭から, 速度 20 の速さで前後の実測角を結ぶ直線 (関節空間. 全関節が同時に着く) を進み,
#             着いたら抜けの終わりまで止まる. 速さは trim2 の元の記録 (run_pick_place 19 本. 速度 20 で
#             0.32 秒おきに動いている間も読めていた) の, 動いている途中の先頭の関節の速さ 128 区間の中央値 29.8°/秒
GAP_FILL_MODES = ("linear", "speed20")
GAP_SEC = 0.25                  # 実測角の間がこれより長いところを「抜け」とみなす (ふだんは 0.1 秒おき)
SPEED20_DEG_PER_SEC = 30.0      # 速度 20 で動いている間の先頭の関節の速さ (上の 29.8°/秒)
GAP_MOVE_MIN_DEG = 0.5          # 抜けの前後でこれより動いていれば「動いた抜け」として数える
# ---- 着いたあとの止まりの詰め方 (--trim-hold. 2026-10-31 依頼 48. corr_v3) ----
# 推論・補正では 1 回の send_angles ごとに daemon の読み直し (約 1.6 秒) を待つので, 中継点・補正の 1 刻みの
# たびに「動いて止まる」をくり返す. trim2 の決まり (止まりの区間ごとに先頭 2 コマ) では, 区間が 2〜3 コマ
# おきに来るため止まりが約 50% 残る (corr_v2). この止まりは送信の確かめを待つだけで, 動作の意味はない.
#   model の区間       : 残さない (中継点の動きだけをつなぐ)
#   correct の区間     : 先頭 1 コマ (補正のコマを減らしすぎない. 補正は 1 刻みが 1 コマで動き終わる)
#   爪の切り替えの直前・直後の区間 : trim2 と同じ先頭 2 コマ (爪が動き終わるのを待つ意味がある) + gripper_guard
HOLD_KEEP_MODEL = 0
HOLD_KEEP_CORRECT = 1
HOLD_KEEP_GRIP = bde.MID_IDLE_KEEP_STEPS


def _rmtree(path):
    """exFAT では本体を消すと macOS が ._ (AppleDouble) も一緒に消すため, 後から ._ を消そうとして
    FileNotFoundError になる. もう無いものは飛ばして消す."""
    def _skip_missing(func, p, exc):
        exc = exc[1] if isinstance(exc, tuple) else exc   # onerror (3.11 まで) は exc_info
        if not isinstance(exc, FileNotFoundError):
            raise exc
    if sys.version_info >= (3, 12):
        shutil.rmtree(path, onexc=_skip_missing)
    else:
        shutil.rmtree(path, onerror=_skip_missing)


def load_stage_times(path):
    """states.csv の stage 行 (状態の切り替わり) と, 各行の stage 列."""
    ts, modes = [], []
    with path.open(newline="") as f:
        for row in csv.DictReader(f):
            if row.get("event_type") != "stage":
                continue
            try:
                ts.append(float(row["ts"])); modes.append(row.get("stage") or MODE_MODEL)
            except Exception:
                continue
    return np.array(ts), modes


def load_log(path):
    out = []
    if not path.exists():
        return out
    for line in path.read_text().splitlines():
        line = line.strip()
        if line:
            try:
                out.append(json.loads(line))
            except Exception:
                continue
    return out


def idle_keep_mask(state, actions):
    """build_demo_episodes.build_one の「停止区間を詰める」(IDLE_*) と同じ決まり. 戻り値: (keep, 詰めた数)."""
    n = len(state)
    motion = np.max(np.abs(actions - state), axis=1)
    keep = np.ones(n, dtype=bool)
    i, trimmed = 0, 0
    while i < n:
        if motion[i] < bde.IDLE_EPS_DEG:
            j = i
            while j < n and motion[j] < bde.IDLE_EPS_DEG:
                j += 1
            if (j - i) / FPS > bde.IDLE_MAX_SEC:
                cut_from = i + int(bde.IDLE_KEEP_SEC * FPS)
                keep[cut_from:j] = False
                trimmed += j - cut_from
            i = j
        else:
            i += 1
    guard = bde.gripper_guard_mask(state)
    trimmed -= int((~keep & guard).sum())
    keep |= guard
    return keep, trimmed


def fill_motion_gaps(t, at, aq, mode="speed20", gap_sec=GAP_SEC, speed=SPEED20_DEG_PER_SEC):
    """10Hz の時刻 t の関節角を, 実測角 (at, aq) から作る. 実測角があるところは実測角の線形補間.

    実測角の間が gap_sec より長い「抜け」[ta, tb] だけ, mode で埋め方を変える:
      linear  : qa → qb を ta〜tb でまっすぐ結ぶ (np.interp と同じ)
      speed20 : ta から qa → qb の直線を, いちばん大きく動く関節が speed (°/秒) になる速さで進み,
                着いたら tb まで qb で止まる. 着くまでの時間が抜けより長いときは抜けいっぱいを使う (linear と同じ)
    戻り値: (q (len(t), 6), 抜けの一覧 [{ta, tb, move_deg, move_sec}])
    """
    if mode not in GAP_FILL_MODES:
        raise ValueError(f"gap fill mode: {mode}")
    q = np.stack([np.interp(t, at, aq[:, k]) for k in range(6)], axis=1)
    gaps = []
    for i in np.flatnonzero(np.diff(at) > gap_sec):
        ta, tb, qa, qb = at[i], at[i + 1], aq[i], aq[i + 1]
        d = float(np.max(np.abs(qb - qa)))
        dur = min(d / speed, tb - ta) if mode == "speed20" else tb - ta
        gaps.append({"ta": float(ta), "tb": float(tb), "move_deg": d, "move_sec": float(dur)})
        if mode == "linear":
            continue
        m = (t > ta) & (t < tb)
        if not m.any():
            continue
        frac = np.ones(int(m.sum())) if dur <= 0 else np.clip((t[m] - ta) / dur, 0.0, 1.0)
        q[m] = qa + frac[:, None] * (qb - qa)
    return q, gaps


def logged_sends(log):
    """log.jsonl から, 送った目標とそのときの時刻の手がかりを集める.

    chunk  : 推論のチャンク. 実行した窓 [t_exec_start, t_exec_end] と, 送った順の目標 sent_targets
    correct: 補正の 1 刻み (送ったもの). 送る直前の時刻 t と目標 after
    """
    chunks, corr = [], []
    for r in log:
        if r.get("type") is None and "chunk" in r and r.get("t_exec_start") is not None:
            chunks.append({"chunk": r["chunk"], "t0": float(r["t_exec_start"]),
                           "t1": float(r.get("t_exec_end") or r["t_exec_start"]),
                           "targets": [np.asarray(x, dtype=float) for x in (r.get("sent_targets") or [])]})
        elif r.get("type") == "correction" and r.get("action") == "move" and r.get("sent"):
            corr.append({"t": float(r["t"]), "target": np.asarray(r["after"], dtype=float)})
    return chunks, corr


def match_gaps_to_sends(gaps, log, min_move=GAP_MOVE_MIN_DEG):
    """動いた抜けが, 記録にある送信 (チャンクの中継点 / 補正の 1 刻み) のどれに当たるかを数える.

    補正: 送る直前の時刻 t が [ta - 0.3, tb] に入る抜け. チャンク: 実行した窓の中の動いた抜けを
    送った順に sent_targets と組にする (数が合うときだけ). 組にできたものは, 抜けの終わりの実測角と
    目標の差 (いちばん大きい関節) を出す. 埋め方そのものには使わない (確かめるためだけ).
    """
    chunks, corr = logged_sends(log)
    moved = [g for g in gaps if g["move_deg"] > min_move]
    out = {"motion_gaps": len(moved), "matched_correction": 0, "matched_chunk": 0,
           "chunk_count_mismatch": 0, "ambiguous": 0, "unmatched": 0, "max_arrival_err_deg": None}
    errs, used = [], set()
    # codex 推奨: 候補の抜けが 2 つ以上ある補正 / 1 つの抜けに補正が 2 つ以上入るものは ambiguous として
    # 組にせず, 差の集計から外す
    cand = [[j for j, g in enumerate(moved) if g["ta"] - 0.3 <= c["t"] <= g["tb"]] for c in corr]
    hits = {}
    for js in cand:
        for j in js:
            hits[j] = hits.get(j, 0) + 1
    for c, js in zip(corr, cand):
        if len(js) == 1 and hits[js[0]] == 1:
            used.add(js[0]); out["matched_correction"] += 1
            errs.append(float(np.max(np.abs(moved[js[0]]["qb"] - c["target"]))))
        elif js:
            out["ambiguous"] += 1
            used.update(js)
    for ch in chunks:
        inside = [j for j, g in enumerate(moved) if j not in used and ch["t0"] - 0.3 <= g["ta"] <= ch["t1"]]
        if len(inside) != len(ch["targets"]):
            out["chunk_count_mismatch"] += 1
            continue
        for j, tg in zip(inside, ch["targets"]):
            used.add(j); out["matched_chunk"] += 1
            errs.append(float(np.max(np.abs(moved[j]["qb"] - tg))))
    out["unmatched"] = len(moved) - len(used)
    out["note"] = "確かめのためだけ (埋め方には使わない). チャンクは数が合うとき送った順に組にする"
    if errs:
        out["max_arrival_err_deg"] = round(max(errs), 2)
    return out


def next_state_actions(state):
    """actions = 次のステップの state (最後は自分自身). 7 列とも."""
    return np.concatenate([state[1:], state[-1:]], axis=0)


def hold_keep_mask(state, actions, modes, keep_model=HOLD_KEEP_MODEL, keep_correct=HOLD_KEEP_CORRECT,
                   keep_grip=HOLD_KEEP_GRIP):
    """bde.mid_idle_keep_mask と同じ止まりの判定・同じずれの切り方で, 区間ごとに残す数を変える (依頼 48).

    止まりの区間 (つながった止まりのコマ) ごとに残す数を決める:
      爪の切り替えのコマ (action[6] != state[6]) のすぐ前で終わる / すぐ後で始まる区間 → keep_grip
      それ以外は区間の先頭のコマの状態で, correct → keep_correct / それ以外 → keep_model
    区間の中で先頭からのずれが MID_IDLE_DRIFT_DEG を超えたら, そこから数え直す (mid_idle_keep_mask と同じ).
    codex 必須 (依頼 48): 数え直すコマ (ずれの切れ目) は残す数が 0 でも残す. 0.2° ずつ動き続ける区間が
    1 コマの大きな跳びにならないように (ゆっくりした動きは 0.5° ごとの姿勢が残る).
    爪の切り替えの前後 (gripper_guard_mask) とエピソードの両端は必ず残す.
    keep_model = keep_correct = keep_grip = MID_IDLE_KEEP_STEPS なら mid_idle_keep_mask と同じ結果.
    """
    n = len(state)
    joint_still = np.max(np.abs(actions[:, :6] - state[:, :6]), axis=1) < bde.MID_IDLE_EPS_DEG
    grip_step = actions[:, 6] != state[:, 6]
    idle = joint_still & ~grip_step
    keep = np.ones(n, dtype=bool)
    i = 0
    while i < n:
        if not idle[i]:
            i += 1
            continue
        j = i
        while j < n and idle[j]:
            j += 1
        if (j < n and grip_step[j]) or (i > 0 and grip_step[i - 1]):
            k = keep_grip
        else:
            k = keep_correct if modes[i] == MODE_CORRECT else keep_model
        anchor, kept = i, 0
        for x in range(i, j):
            if np.max(np.abs(state[x, :6] - state[anchor, :6])) > bde.MID_IDLE_DRIFT_DEG:
                anchor, kept = x, 0
                if k == 0:
                    continue        # ずれの切れ目は残す (keep[x] は True のまま)
            if kept < k:
                kept += 1
            else:
                keep[x] = False
        i = j
    keep |= bde.gripper_guard_mask(state)
    if n:
        keep[0] = True
        keep[-1] = True
    return keep


def select_steps(t, q, g, mode_ts, modes, windows, trim_mid_idle=True, trim_hold=False):
    """学習に使うステップを選び, state / actions を作る (ファイルを読まない部分. テストで使う).

    戻り値: dict(idx=元の t の添字, state, actions, modes, counts)
    """
    state_all = np.concatenate([q, g[:, None]], axis=1).astype(np.float32)
    keep, m_all = step_keep_mask(t, mode_ts, modes, windows)
    idx = np.flatnonzero(keep)
    counts = {"grid_steps": int(len(t)),
              "excluded_mode": int((~np.isin(m_all, (MODE_MODEL, MODE_CORRECT))).sum()),
              "excluded_window": int((np.isin(m_all, (MODE_MODEL, MODE_CORRECT)) & ~keep).sum())}
    if len(idx) == 0:
        return {"idx": idx, "state": state_all[:0], "actions": state_all[:0], "modes": [], "counts": counts}
    state = state_all[idx]
    # 除いたあとに残ったステップでつなぐ. actions は「つないだ並びでの次の state」
    actions = next_state_actions(state)
    k1, idle_trimmed = idle_keep_mask(state, actions)
    idx, state, actions = idx[k1], state[k1], actions[k1]
    counts["idle_trimmed"] = int(idle_trimmed)
    counts["steps_before_mid_trim"] = int(len(idx))
    mid = 0
    if trim_mid_idle and len(idx):
        if trim_hold:
            mk = hold_keep_mask(state, actions, [str(m_all[k]) for k in idx])
        else:
            mk = bde.mid_idle_keep_mask(state, actions)
        mid = int((~mk).sum())
        idx, state, actions = idx[mk], state[mk], actions[mk]
    # trim2 と同じ: 関節 6 列は残ったステップで作り直す. グリッパー列は元の (次の) 指令のまま
    actions = actions.copy()
    actions[:-1, :6] = state[1:, :6]
    actions[-1, :6] = state[-1, :6]
    counts["mid_idle_trimmed"] = mid
    sel_modes = [str(m_all[k]) for k in idx]
    counts["model_steps"] = sel_modes.count(MODE_MODEL)
    counts["correct_steps"] = sel_modes.count(MODE_CORRECT)
    return {"idx": idx, "state": state, "actions": actions, "modes": sel_modes, "counts": counts}


def build_one(run_dir, out_dir, args):
    stem = run_dir.name
    trial_p = run_dir / f"{stem}.json"
    rec = run_dir / "recording"
    if not trial_p.exists():
        return {"run": stem, "used": False, "reason": f"missing:{trial_p.name}"}
    trial = json.loads(trial_p.read_text())
    if trial.get("recording_enabled") is not True:
        return {"run": stem, "used": False, "reason": "no_recording"}
    if trial.get("outcome") != "success":
        return {"run": stem, "used": False, "reason": f"outcome={trial.get('outcome')}"}
    need = [rec / "states.csv", rec / "top.mp4", rec / "top_frame_log.csv"]
    missing = [p.name for p in need if not p.exists()]
    if missing:
        return {"run": stem, "used": False, "reason": "missing:" + ",".join(missing)}
    log = load_log(run_dir / "log.jsonl")
    unfinished = [r for r in log if r.get("type") == "rewind" and not r.get("finished", True)]
    unfinished += [r for r in log if r.get("type") == "correction" and r.get("action") == "cancel"
                   and not r.get("finished", True)]
    if unfinished:
        # 戻りきれなかった巻き戻しがあると, つないだ前後で姿勢が続かない
        return {"run": stem, "used": False, "reason": f"rewind_unfinished x{len(unfinished)}"}

    at, aq, gt, gs = bde.load_states(rec / "states.csv")
    if len(at) < bde.MIN_ACTUAL_SAMPLES:
        return {"run": stem, "used": False, "reason": f"actual 実測が少ない ({len(at)})"}
    order = np.argsort(at); at, aq = at[order], aq[order]
    at, aq = bde.drop_spikes(at, aq)
    try:
        tm = json.loads((rec / "top_meta.json").read_text())
    except Exception:
        tm = {}
    if tm.get("status") != "ok":
        return {"run": stem, "used": False, "reason": f"recording_status={tm.get('status')}"}
    fidx, fts = bde.load_frame_log(rec / "top_frame_log.csv")
    if len(fts) < FPS:
        return {"run": stem, "used": False, "reason": f"top フレームが少ない ({len(fts)})"}
    t0, t1 = max(fts[0], at[0]), min(fts[-1], at[-1])
    if t1 - t0 < 3.0:
        return {"run": stem, "used": False, "reason": "記録の重なりが短い"}
    t = np.arange(t0, t1, 1.0 / FPS)
    gap_fill = getattr(args, "gap_fill", "linear")
    q, gaps = fill_motion_gaps(t, at, aq, mode=gap_fill)
    for g in gaps:
        k = int(np.searchsorted(at, g["tb"]))
        g["qb"] = aq[k]
    gap_match = match_gaps_to_sends(gaps, log)
    g = np.zeros(len(t))
    if len(gt):
        go = np.argsort(gt); gt, gs = gt[go], gs[go]
        for i, ti in enumerate(t):
            k = np.searchsorted(gt, ti, side="right") - 1
            g[i] = gs[k] if k >= 0 else 0.0

    mode_ts, modes = load_stage_times(rec / "states.csv")
    windows = excluded_windows(log)
    trim_hold = getattr(args, "trim_hold", False)
    sel = select_steps(t, q, g, mode_ts, modes, windows, trim_mid_idle=not args.no_trim_mid_idle,
                       trim_hold=trim_hold)
    if len(sel["idx"]) < FPS:
        return {"run": stem, "used": False, "reason": f"使えるステップが少ない ({len(sel['idx'])})"}
    tt = t[sel["idx"]]

    # codex R1-4: 画像の候補も同じ決まりで絞る (除いた区間・rewind / home の画像を, 境界の
    # 最近傍として拾わない). 候補から選ぶので, 除いた時刻の画像は入らない
    fkeep, _ = step_keep_mask(fts, mode_ts, modes, windows)
    fidx, fts = fidx[fkeep], fts[fkeep]
    if len(fts) < 2:
        return {"run": stem, "used": False, "reason": "使える画像が少ない"}
    nearest = np.clip(np.searchsorted(fts, tt), 1, len(fts) - 1)
    left = fts[nearest - 1]; right = fts[nearest]
    nearest = np.where(np.abs(tt - left) <= np.abs(right - tt), nearest - 1, nearest)
    want = {}
    for step, n in enumerate(nearest):
        want.setdefault(int(fidx[n]), []).append(step)

    if out_dir.exists():
        if not args.force:
            return {"run": stem, "used": True, "reason": "既存 (スキップ・--force で再生成)",
                    "steps": int(len(tt))}
        _rmtree(out_dir)
    (out_dir / "frames").mkdir(parents=True)
    cap = cv2.VideoCapture(str(rec / "top.mp4"))
    fi, written, max_need, size = 0, 0, max(want), None
    while fi <= max_need:
        ok, frame = cap.read()
        if not ok:
            break
        if fi in want:
            h, w = frame.shape[:2]
            s = bde.IMG_LONG / max(h, w)
            small = cv2.resize(frame, (round(w * s), round(h * s)), interpolation=cv2.INTER_AREA)
            size = small.shape[1], small.shape[0]
            for step in want[fi]:
                # codex R3 P2: 保存に失敗したコマは数えない (欠けたエピソードを採用しない)
                if cv2.imwrite(str(out_dir / "frames" / f"{step:06d}.jpg"), small,
                               [cv2.IMWRITE_JPEG_QUALITY, bde.JPEG_QUALITY]):
                    written += 1
        fi += 1
    cap.release()
    if written != len(tt):
        _rmtree(out_dir)
        return {"run": stem, "used": False,
                "reason": f"フレーム抽出不足 written={written} steps={len(tt)} (mp4 と frame_log の不一致)"}

    np.savez_compressed(out_dir / "episode.npz",
                        t=(np.arange(len(tt)) / FPS).astype(np.float64),
                        t_orig=tt - tt[0], state=sel["state"], actions=sel["actions"])
    c = sel["counts"]
    meta = {
        "run_stem": stem, "task": bde.TASK, "fps": FPS, "steps": int(len(tt)),
        "duration_sec": round(len(tt) / FPS, 1),
        "idle_trimmed_sec": round(c.get("idle_trimmed", 0) / FPS, 1),
        "trim_mid_idle": not args.no_trim_mid_idle,
        "trim_hold": ({"keep_model": HOLD_KEEP_MODEL, "keep_correct": HOLD_KEEP_CORRECT,
                       "keep_grip": HOLD_KEEP_GRIP} if trim_hold and not args.no_trim_mid_idle else None),
        "steps_before_mid_trim": c.get("steps_before_mid_trim"),
        "mid_idle_trimmed_steps": c.get("mid_idle_trimmed", 0),
        "image_size": size, "state_dim": 7,
        "state_names": ["j0", "j1", "j2", "j3", "j4", "j5", "gripper"],
        "action": "absolute target = state 0.1s later",
        "actual_samples": int(len(at)), "gripper_cmds": int(len(gt)),
        "human_verdict": {"pick": "yes", "place": "yes"},
        "source": "pi05_policy (model + correct)",
        "model_steps": c.get("model_steps", 0), "correct_steps": c.get("correct_steps", 0),
        "excluded_mode_steps": c["excluded_mode"], "excluded_window_steps": c["excluded_window"],
        "excluded_windows": [[round(a, 3), round(b, 3)] for a, b in windows],
        "rewinds": sum(1 for r in log if r.get("type") == "rewind"),
        "corrections_accepted": sum(1 for r in log if r.get("type") == "correction"
                                    and r.get("action") == "accept"),
        # 依頼 46: 補正の数は「補正の操作として記録された行の数」(log.jsonl の type=correction の行. move / accept /
        # cancel / gripper など全部) で数える
        "correction_rows": sum(1 for r in log if r.get("type") == "correction"),
        "correction_moves": sum(1 for r in log if r.get("type") == "correction" and r.get("action") == "move"),
        "gap_fill": {"mode": gap_fill, "gap_sec": GAP_SEC,
                     "speed_deg_per_sec": SPEED20_DEG_PER_SEC if gap_fill == "speed20" else None,
                     "gaps": len(gaps), **gap_match},
        "max_step_deg": [round(float(v), 2) for v in
                         (np.abs(np.diff(sel["state"][:, :6], axis=0)).max(axis=0) if len(sel["state"]) > 1
                          else np.zeros(6))],
        "step_modes": sel["modes"],
    }
    (out_dir / "meta.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False))
    return {"run": stem, "used": True, "reason": "ok", "steps": int(len(tt)),
            "model_steps": meta["model_steps"], "correct_steps": meta["correct_steps"],
            "rewinds": meta["rewinds"], "corrections_accepted": meta["corrections_accepted"],
            "correction_rows": meta["correction_rows"], "max_step_deg": meta["max_step_deg"]}


def check_compat(ours, ref):
    """補正エピソードがお手本と一緒に学習に使えるか (形・キー・FPS・画像の大きさ). 問題の一覧を返す."""
    probs = []

    def _eps(d):
        return sorted(p.parent for p in Path(d).glob("*/episode.npz"))
    a, b = _eps(ours), _eps(ref)
    if not a:
        return ["補正エピソードがない"]
    if not b:
        return ["お手本エピソードがない"]
    rm = json.loads((b[0] / "meta.json").read_text())
    rz = np.load(b[0] / "episode.npz")
    for ep in a:
        m = json.loads((ep / "meta.json").read_text())
        z = np.load(ep / "episode.npz")
        if sorted(z.files) != sorted(rz.files):
            probs.append(f"{ep.name}: npz のキー {sorted(z.files)} != お手本 {sorted(rz.files)}")
        for k in ("state", "actions"):
            if z[k].dtype != rz[k].dtype or z[k].shape[1:] != rz[k].shape[1:]:
                probs.append(f"{ep.name}: {k} {z[k].dtype}{z[k].shape} != お手本 {rz[k].dtype}{rz[k].shape}")
        for k in ("fps", "image_size", "state_dim", "state_names", "task"):
            if m.get(k) != rm.get(k):
                probs.append(f"{ep.name}: meta {k}={m.get(k)} != お手本 {rm.get(k)}")
        n = len(z["state"])
        # exFAT に macOS が作る AppleDouble (._000000.jpg) は画像ではないので数えない
        nf = len(list((ep / "frames").glob("[!.]*.jpg")))
        if not (n == len(z["actions"]) == nf == len(z["t"])):
            probs.append(f"{ep.name}: 長さが合わない state={n} actions={len(z['actions'])} frames={nf}")
        img = cv2.imread(str(ep / "frames" / "000000.jpg"))
        if img is None or (img.shape[1], img.shape[0]) != tuple(rm["image_size"]):
            probs.append(f"{ep.name}: 画像の大きさ {None if img is None else img.shape[:2][::-1]} "
                         f"!= {rm['image_size']}")
    return probs


def write_readme(ds, results, args):
    used = [r for r in results.values() if r.get("used")]
    lines = [
        f"# {ds.name}", "",
        "実機推論 (`pi05/run_policy_client.py --execute`) で, 失敗しかけたら巻き戻し (b) /",
        "キーボードで関節を直し (c) ながら成功させた run から作ったエピソード。",
        "お手本 `mecharm_redblock_box_teach_v1_trim2` と同じ形 (frames/ + episode.npz + meta.json, 10Hz)。", "",
        "- 作成: `python3 pi05/build_correction_episodes.py`"
        + ("" if not args.no_trim_mid_idle else " --no-trim-mid-idle")
        + f" --gap-fill {args.gap_fill}"
        + (" --trim-hold" if getattr(args, "trim_hold", False) else ""),
        f"- 元データ: `runs/pi05_policy/<run_stem>/` (読むだけ。上書き・移動・削除しない)", "",
        "## 作り方", "",
        "- 結果が「成功」(y) の run だけ使う (`{run_stem}.json` の outcome)",
        "- 状態が model / correct のステップだけ使う。rewind / home / wait (確認の入力待ち。手が写り得る) / "
        "regrip (補正の R で開いてから閉じ終わるまで) は除く",
        "- 巻き戻しで取り消した区間 (区間を始めた時刻 〜 巻き戻しが終わった時刻) と, 取り消した補正は除き,",
        "  その手前と巻き戻したあとの続きをつなぐ",
        "- actions は「0.1 秒後の state」(つないだ並びで作り直す。グリッパー列は次の指令)",
        "- 止まっている区間は build_demo_episodes と同じ決まり (IDLE_* と mid_idle_keep_mask) で詰める"
        + ("" if not getattr(args, "trim_hold", False) else
           f"。ただし途中の止まりは区間ごとに残す数を変える (--trim-hold): model {HOLD_KEEP_MODEL} コマ / "
           f"correct {HOLD_KEEP_CORRECT} コマ / 爪の切り替えの直前・直後 {HOLD_KEEP_GRIP} コマ (trim2 と同じ) + 切り替えの前後 1 コマ"),
        "- 実測角の抜け (送信の間 daemon とのつながりがふさがる. 約 1.6 秒) の埋め方: "
        + ("前後の実測角をまっすぐ結ぶ (linear)" if args.gap_fill == "linear" else
           f"抜けの頭から速度 20 の速さ ({SPEED20_DEG_PER_SEC:.0f}°/秒. いちばん大きく動く関節) で前後の実測角を結ぶ"
           "直線を進み, 着いたら止まる (speed20)"),
        "- 戻りきれなかった巻き戻しがある run は使わない (前後の姿勢が続かない)", "",
        "## 数", "",
        f"採用 {len(used)} 本 / 総ステップ {sum(r.get('steps', 0) for r in used)} "
        f"(model {sum(r.get('model_steps', 0) for r in used)} / correct {sum(r.get('correct_steps', 0) for r in used)})", "",
        "| run_stem | 採用 | steps | model | correct | 巻き戻し | 補正 (受け入れ) | 補正の行 | 1 コマの最大 (j0〜j5) | 理由 |",
        "| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |",
    ]
    for k in sorted(results):
        r = results[k]
        lines.append(f"| {k} | {'○' if r.get('used') else '-'} | {r.get('steps', '')} | "
                     f"{r.get('model_steps', '')} | {r.get('correct_steps', '')} | {r.get('rewinds', '')} | "
                     f"{r.get('corrections_accepted', '')} | {r.get('correction_rows', '')} | "
                     f"{' '.join(str(v) for v in r.get('max_step_deg') or [])} | {r.get('reason', '')} |")
    (ds / "README.md").write_text("\n".join(lines) + "\n")


def copy_code(ds):
    """データと対のコード: このデータを作ったスクリプト (と読み込む部品) を ds/code/ に写す."""
    dest = ds / "code"
    dest.mkdir(parents=True, exist_ok=True)
    for name in ("build_correction_episodes.py", "rewind_correct.py", "build_demo_episodes.py"):
        shutil.copy2(_HERE / name, dest / name)
    return dest


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dataset", default=DEFAULT_DATASET)
    ap.add_argument("--runs-root", default=str(RUNS_ROOT), help="run_policy_client の記録の置き場所")
    ap.add_argument("--run", action="append", help="特定の run_stem だけ (複数可)")
    ap.add_argument("--force", action="store_true", help="既存エピソードを作り直す")
    ap.add_argument("--no-trim-mid-idle", action="store_true",
                    help="途中の停止を mid_idle_keep_mask で詰めない (既定は詰める = trim2 と同じ)")
    ap.add_argument("--check-with", help="お手本の raw_episodes. 一緒に学習に使えるかを確かめる")
    ap.add_argument("--gap-fill", choices=GAP_FILL_MODES, default="linear",
                    help="実測角の抜けの埋め方. linear = 前後をまっすぐ (corr_v1) / speed20 = 速度 20 の速さで進んで止まる "
                         "(corr_v2. 依頼 46)")
    ap.add_argument("--trim-hold", action="store_true",
                    help="中継点・補正の 1 刻みごとの着いたあとの止まりを区間ごとに詰める (corr_v3. 依頼 48)")
    ap.add_argument("--data-root", help=argparse.SUPPRESS)   # テスト用 (KIOXIA の代わり)
    args = ap.parse_args()
    if args.trim_hold and args.no_trim_mid_idle:
        sys.exit("[ERROR] --trim-hold は途中の止まりを詰めるときの決まり. --no-trim-mid-idle とは一緒に使えない")
    if args.data_root:
        root = Path(args.data_root)
    else:
        bde.ensure_ssd()
        root = DATA_ROOT
    runs_root = Path(args.runs_root)
    ds = root / "datasets" / args.dataset
    out_root = ds / "raw_episodes"
    # 既存のお手本データセット (pi05_policy 以外から作ったエピソード) には書かない (codex R1 B3)
    others = [p.parent.name for p in out_root.glob("*/meta.json")
              if not str(json.loads(p.read_text()).get("source", "")).startswith("pi05_policy")]
    if others:
        sys.exit(f"[ABORT] 出力先に補正エピソード以外が入っています: {ds} (例: {', '.join(others[:3])})\n"
                 "        元データを上書きしないよう, --dataset に別の名前を指定してください")
    runs = sorted(p for p in runs_root.glob("pi05_*") if p.is_dir() and (p / "recording").is_dir())
    if args.run:
        runs = [p for p in runs if p.name in set(args.run)]
    if not runs:
        print(f"[INFO] 記録つきの run がない: {runs_root}")
    out_root.mkdir(parents=True, exist_ok=True)
    results = {}
    for r in runs:
        try:
            res = build_one(r, out_root / r.name, args)
        except Exception as e:
            res = {"run": r.name, "used": False, "reason": f"exception:{type(e).__name__}:{e}"}
        results[r.name] = res
        print(f"[{'OK ' if res['used'] else '-- '}] {res['run']}  {res.get('steps', '')}  {res['reason']}")
    index_path = ds / "index.json"
    old = json.loads(index_path.read_text()) if index_path.exists() else {}
    old.update(results)
    index_path.write_text(json.dumps(old, indent=2, ensure_ascii=False))
    write_readme(ds, old, args)
    print(f"[CODE] データと対のコード → {copy_code(ds)}")
    used = [v for v in old.values() if v.get("used")]
    print(f"\n採用エピソード合計: {len(used)}  総ステップ: {sum(v.get('steps', 0) for v in used)}  → {out_root}")
    if args.check_with:
        probs = check_compat(out_root, args.check_with)
        print("[CHECK] お手本と一緒に使えるか: " + ("OK" if not probs else f"NG {len(probs)} 件"))
        for p in probs:
            print("        " + p)
        return 1 if probs else 0
    return 0


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