#!/usr/bin/env python3
"""diag_fill_flags.py — 既存の補正データセットで, 実測角の抜けを埋めたコマの割合を測る (2026-11-14 依頼 196).

datasets/ の名前に corr / deep_even10 が入るデータセットの, 補正エピソード (meta.json の source が pi05_policy)
ごとに, 元の runs/pi05_policy/{run_stem}/ から build_correction_episodes.build_one を meta.json と同じ条件
(gap_fill の mode・trim_mid_idle・trim_hold) で /tmp に作り直し, state_source (コマごとの印) を数える.
作り直した state / actions / t がデータセットの episode.npz と一致しない本は数えない (理由を CSV に残す).

データセット・元の記録は読むだけ (書かない). 書くのは /tmp と diag/fill_flags_20261114/ だけ.

  python3 pi05/diag_fill_flags.py
"""
import argparse
import csv
import json
import sys
from pathlib import Path
from types import SimpleNamespace

import numpy as np

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

import build_correction_episodes as bce  # noqa: E402
import build_demo_episodes as bde  # noqa: E402

OUT_NAME = "fill_flags_20261114"
TMP_ROOT = Path("/tmp") / OUT_NAME
# 学習に使ったことがあるデータセットを先に並べる (依頼 196)
TRAINED_FIRST = ("mecharm_redblock_box_teach_v1_trim2_corr_v5", "mecharm_redblock_box_deep_even10_v1",
                 "mecharm_redblock_box_teach_v1_trim2_corr_v5_deep")


def target_datasets(ds_root):
    names = sorted(p.name for p in ds_root.iterdir()
                   if p.is_dir() and not p.name.startswith(".") and ("corr" in p.name or "deep_even10" in p.name))
    return sorted(names, key=lambda n: (n not in TRAINED_FIRST, TRAINED_FIRST.index(n) if n in TRAINED_FIRST else 0, n))


def rebuild(run_stem, meta, cache):
    """meta.json と同じ条件で /tmp に作り直す. 戻り値: (npz の中身 dict / None, 理由)."""
    gf = meta.get("gap_fill") or {}
    key = (run_stem, gf.get("mode"), bool(meta.get("trim_mid_idle")), meta.get("trim_hold") is not None)
    if key in cache:
        return cache[key]
    out = TMP_ROOT / "rebuild" / "_".join(str(k) for k in key) / run_stem
    args = SimpleNamespace(gap_fill=gf.get("mode") or "linear", trim_hold=key[3],
                           no_trim_mid_idle=not key[2], max_protrude_px=None, force=True)
    res = bce.build_one(bce.RUNS_ROOT / run_stem, out, args)
    if not res.get("used"):
        cache[key] = (None, f"作り直せない: {res.get('reason')}")
    else:
        z = np.load(out / "episode.npz")
        cache[key] = ({k: z[k] for k in z.files}, "ok")
    return cache[key]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dataset", action="append", help="このデータセットだけ (複数可)")
    args = ap.parse_args()
    bde.ensure_ssd()
    ds_root = bde.DATA_ROOT / "datasets"
    out_dir = bde.DATA_ROOT / "diag" / OUT_NAME
    if out_dir.exists():
        sys.exit(f"[ABORT] 出力先がもうある: {out_dir}")
    names = args.dataset or target_datasets(ds_root)
    rows, cache = [], {}
    for name in names:
        for mp in sorted((ds_root / name / "raw_episodes").glob("[!.]*/meta.json")):
            meta = json.loads(mp.read_text())
            if not str(meta.get("source", "")).startswith("pi05_policy"):
                continue
            ep = mp.parent
            stem = meta.get("run_stem", ep.name)
            z = np.load(ep / "episode.npz")
            row = {"dataset": name, "episode": ep.name, "run_stem": stem,
                   "gap_fill": (meta.get("gap_fill") or {}).get("mode"),
                   "trim_hold": meta.get("trim_hold") is not None, "steps": len(z["state"]),
                   "filled_steps": "", "filled_motion_steps": "", "filled_ratio": "",
                   "gaps": (meta.get("gap_fill") or {}).get("gaps"),
                   "motion_gaps": (meta.get("gap_fill") or {}).get("motion_gaps"), "status": ""}
            new, why = rebuild(stem, meta, cache)
            if new is None:
                row["status"] = why
            elif not all(k in new and new[k].shape == z[k].shape and new[k].tobytes() == z[k].tobytes()
                         for k in ("state", "actions", "t")):
                row["status"] = "state が一致しない (数えない)"
            else:
                src = new["state_source"]
                n_f = int((src != bce.SRC_MEASURED).sum())
                row.update(filled_steps=n_f, filled_motion_steps=int((src == bce.SRC_FILLED_MOVE).sum()),
                           filled_ratio=round(n_f / len(src), 4), status="ok")
            rows.append(row)
            print(f"{name}  {ep.name}  {row['status']}  {row['filled_steps']}/{row['steps']}")

    if not rows:   # codex P2: 出力先を作る前に止める (空のまま作ると次に走らせられない)
        sys.exit("[ABORT] 対象の補正エピソードがない")
    out_dir.mkdir(parents=True)
    with (out_dir / "episodes.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader(); w.writerows(rows)
    summ = []
    for name in names:
        rs = [r for r in rows if r["dataset"] == name]
        ok = [r for r in rs if r["status"] == "ok"]
        steps = sum(r["steps"] for r in ok); filled = sum(r["filled_steps"] for r in ok)
        top = max(ok, key=lambda r: r["filled_ratio"]) if ok else None
        summ.append({"dataset": name, "corr_episodes": len(rs), "counted": len(ok),
                     "not_counted": len(rs) - len(ok), "steps": steps, "filled_steps": filled,
                     "filled_motion_steps": sum(r["filled_motion_steps"] for r in ok),
                     "filled_ratio": round(filled / steps, 4) if steps else "",
                     "max_episode": top["episode"] if top else "",
                     "max_episode_ratio": top["filled_ratio"] if top else ""})
    with (out_dir / "summary.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(summ[0].keys()))
        w.writeheader(); w.writerows(summ)
    (out_dir / "README.md").write_text(
        "# fill_flags_20261114 (依頼 196)\n\n"
        "既存の補正データセットで, 実測角の抜けを埋めたコマ (state_source が 1 または 2) の数と割合。\n"
        "`python3 pi05/diag_fill_flags.py` で作った。データセットは読むだけ。作り直しは /tmp。\n\n"
        "- state_source: 0 = 実測角 / 1 = 動いた抜けを埋めた値 / 2 = 動かなかった抜けを埋めた値\n"
        "- filled_ratio = (1 と 2 のコマ) / 全コマ。filled_motion_steps は 1 のコマの数\n"
        "- status が ok でない本 (作り直した state / actions / t がデータセットと一致しない・作り直せない) は数えない\n"
        "- episodes.csv: 本ごと / summary.csv: データセットごと (max_episode = 割合がいちばん多い本)\n")
    print(f"→ {out_dir}")
    return 0


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