#!/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

作り方:
  - 使うのは 結果が「成功」(y) の run だけ
  - 時間軸は 10Hz. state は実測角の線形補間 + グリッパーは指令の時刻で 0/1 (build_demo_episodes と同じ)
  - 学習に使うのは状態が model / correct のステップだけ. rewind / home / wait (確認の入力待ち.
    人がブロックを置き直すと手が写る) のステップと,
    巻き戻しで取り消した区間 (区間を始めた時刻 〜 巻き戻しが終わった時刻) と,
    取り消した補正 (補正を始めた時刻 〜 戻り終えた時刻) は除く.
    除いたあとは「その手前」と「巻き戻したあとの続き」をつなぐ (戻った先は区間の始めの姿勢なので続く)
  - actions は「0.1 秒後の state」. 除いたあとに残ったステップで作り直す (trim2 と同じ作り方)
  - 止まっている区間は build_demo_episodes と同じ決まりで詰める (IDLE_* と --trim-mid-idle の
    mid_idle_keep_mask. 既定で詰める. --no-trim-mid-idle で前者だけ)
  - 元の記録は読むだけ (上書き・移動・削除しない)

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


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


def select_steps(t, q, g, mode_ts, modes, windows, trim_mid_idle=True):
    """学習に使うステップを選び, 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):
        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)
    q = np.stack([np.interp(t, at, aq[:, k]) for k in range(6)], axis=1)
    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)
    sel = select_steps(t, q, g, mode_ts, modes, windows, trim_mid_idle=not args.no_trim_mid_idle)
    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))}
        shutil.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):
        shutil.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,
        "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"),
        "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"]}


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"])
        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"- 元データ: `runs/pi05_policy/<run_stem>/` (読むだけ。上書き・移動・削除しない)", "",
        "## 作り方", "",
        "- 結果が「成功」(y) の run だけ使う (`{run_stem}.json` の outcome)",
        "- 状態が model / correct のステップだけ使う。rewind / home / wait (確認の入力待ち。手が写り得る) は除く",
        "- 巻き戻しで取り消した区間 (区間を始めた時刻 〜 巻き戻しが終わった時刻) と, 取り消した補正は除き,",
        "  その手前と巻き戻したあとの続きをつなぐ",
        "- actions は「0.1 秒後の state」(つないだ並びで作り直す。グリッパー列は次の指令)",
        "- 止まっている区間は build_demo_episodes と同じ決まり (IDLE_* と mid_idle_keep_mask) で詰める",
        "- 戻りきれなかった巻き戻しがある 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 | 巻き戻し | 補正 | 理由 |",
        "| --- | --- | --- | --- | --- | --- | --- | --- |",
    ]
    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('reason', '')} |")
    (ds / "README.md").write_text("\n".join(lines) + "\n")


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("--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
    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)
    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())