#!/usr/bin/env python3
"""check_dataset.py — 学習に出す前に、お手本データ (raw_episodes) を点検する (Mac 側・読むだけ).

2026-10-21 (lab Log11 / LOG358): teach_v1 で学習したモデルが、ブロックの真上で止まった.
原因はお手本の側で、停止しているステップが 40% あり、真上で 1.5〜2 秒止まっていた.
action_horizon は 10 (1 秒) なので、1 秒を超える停止は「止まる」という動作として丸ごと覚える.
学習の前に数えれば分かったことなので、毎回の学習前にこれを実行する.

点検する項目 (FAIL が 1 つでもあれば終了コード 1):
  1. そろい方   画像の枚数 = ステップ数 / 画像サイズが全エピソードで同じ / meta の task が同じ   違えば FAIL
  2. 停止       停止ステップの割合と、いちばん長い連続         1.0 秒超で FAIL / 割合 20% 超で WARN
  3. グリッパー 切り替えが「閉 → 開」の 2 回か                 違えば WARN
  4. 肩と j4    j1 が J1_SAFE_MAX 超 / j4 が J4_SAFE の外       どちらも FAIL
  5. ばらつき   つかんだ瞬間の j0 を 5° 刻みで数える           1 か所が半分を超えたら WARN
  6. action     |action[t] - state[t+1]| の最大                 3.5° 超で WARN

データには一切書き込まない (np.load と meta.json の読み取りだけ).

Usage:
  python3 pi05/check_dataset.py /Volumes/KIOXIA/r2r2r-data/datasets/mecharm_redblock_box_teach_v1_trim2/raw_episodes
"""
import argparse
import json
import sys
from collections import Counter
from pathlib import Path

import cv2
import numpy as np

# 肩の上限は本線 fk_lookup.py の値をそのまま使う (二重定義しない. record_teach_demo.py と同じ扱い)
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "dynamic_tracking_limits" / "place" / "place_modules"))
from fk_lookup import J1_SAFE_MAX  # noqa: E402  88.0. これを超える指令はアームが黙って捨てる (LOG352)

FPS_DEFAULT = 10                # meta.json の fps が無いときの既定 (build_demo_episodes.py の FPS)
ACTION_HORIZON = 10             # 学習設定の action_horizon (openpi_addon の config と同じ)
STOP_EPS_DEG = 0.3              # build_demo_episodes.py の MID_IDLE_EPS_DEG と同じ「停止」の定義
STOP_RUN_FAIL_SEC = ACTION_HORIZON / FPS_DEFAULT    # 1.0 秒. これを超える連続停止は丸ごと覚える
STOP_RATIO_WARN = 0.20          # 停止ステップの割合がこれを超えたら WARN
J4_SAFE = (-90.0, 60.0)         # run_pick_place.py の J4_SAFE_MIN/MAX
J0_BIN_DEG = 5.0                # つかんだ瞬間の j0 を分ける刻み (eval_policy_offline.py と同じ)
POS_SHARE_WARN = 0.5            # 1 か所がこれを超えたら WARN (LOG348: 同じ場所ばかりで画像を見ずに手を伸ばした)
ACT_STATE_WARN_DEG = 3.5        # |action[t] - state[t+1]| がこれを超えたら WARN
_LEVELS = {"PASS": 0, "WARN": 1, "FAIL": 2}


class Check:
    """1 項目の結果. fail/warn を呼ぶと status がその分だけ悪くなる."""

    def __init__(self, num, title):
        self.num, self.title, self.status, self.lines = num, title, "PASS", []

    def fail(self, msg):
        self._add("FAIL", msg)

    def warn(self, msg):
        self._add("WARN", msg)

    def note(self, msg):
        self.lines.append(("", msg))

    def _add(self, level, msg):
        if _LEVELS[level] > _LEVELS[self.status]:
            self.status = level
        self.lines.append((level, msg))

    def report(self):
        print(f"[{self.status:4s}] {self.num}. {self.title}")
        for level, msg in self.lines:
            print(f"         {('[' + level + '] ') if level else ''}{msg}")


def load_episodes(raw_dir: Path):
    eps = sorted(p.parent for p in raw_dir.glob("*/episode.npz"))
    if not eps:
        sys.exit(f"[ERROR] お手本が見つからない (*/episode.npz): {raw_dir}")
    out = []
    for ep in eps:
        z = np.load(ep / "episode.npz")
        meta_p = ep / "meta.json"
        out.append({
            "name": ep.name,
            "dir": ep,
            "state": np.asarray(z["state"], dtype=np.float64),
            "actions": np.asarray(z["actions"], dtype=np.float64),
            "frames": sorted((ep / "frames").glob("[!.]*.jpg")),   # ._*.jpg (exFAT の AppleDouble) を除く
            "meta": json.loads(meta_p.read_text()) if meta_p.exists() else None,
        })
    return out


def fps_of(e):
    meta = e["meta"] or {}
    try:
        return float(meta.get("fps") or FPS_DEFAULT)
    except (TypeError, ValueError):
        return float(FPS_DEFAULT)


def stop_mask(e):
    """停止ステップ: 6 関節すべて |action - state| < 0.3° かつ グリッパーの action と state が同じ."""
    s, a = e["state"], e["actions"]
    n = min(len(s), len(a))
    if n == 0 or s.shape[1] < 7 or a.shape[1] < 7:
        return np.zeros(0, dtype=bool)
    moved = np.abs(a[:n, :6] - s[:n, :6]) < STOP_EPS_DEG
    return moved.all(axis=1) & (a[:n, 6] == s[:n, 6])


def longest_run(mask):
    best = cur = 0
    for v in mask:
        cur = cur + 1 if v else 0
        best = max(best, cur)
    return best


def grip_changes(e):
    """グリッパーの切り替え: [(ステップ, '閉' or '開'), ...]."""
    g = (e["state"][:, 6] > 0.5).astype(int)
    return [(int(i) + 1, "閉" if g[i + 1] == 1 else "開") for i in np.where(np.diff(g) != 0)[0]]


def check_shape(eps):
    c = Check(1, "そろい方 (画像の枚数・画像サイズ・meta の task)")
    tasks, metas, imgs = Counter(), Counter(), Counter()
    for e in eps:
        s, a, n_img = e["state"], e["actions"], len(e["frames"])
        if s.ndim != 2 or s.shape[1] != 7 or a.shape != s.shape:
            c.fail(f"{e['name']}: state{s.shape} / actions{a.shape} が (N,7) で揃っていない")
            continue
        if n_img != len(s):
            c.fail(f"{e['name']}: 画像 {n_img} 枚 ≠ ステップ {len(s)}")
        if e["meta"] is None:
            c.fail(f"{e['name']}: meta.json がない")
        else:
            tasks[str(e["meta"].get("task"))] += 1
            metas[tuple(e["meta"].get("image_size") or ())] += 1
        if n_img:
            bgr = cv2.imread(str(e["frames"][0]))
            if bgr is None:
                c.fail(f"{e['name']}: 画像を読めない: {e['frames'][0].name}")
            else:
                imgs[(bgr.shape[1], bgr.shape[0])] += 1

    if len(tasks) > 1:
        c.fail("meta の task が揃っていない: " + " / ".join(f"{t!r} x{n}" for t, n in tasks.most_common()))
    elif tasks:
        c.note(f"task {list(tasks)[0]!r} ({sum(tasks.values())} 本すべて同じ)")
    if len(imgs) > 1:
        c.fail("画像サイズが揃っていない: " + " / ".join(f"{w}x{h} x{n}" for (w, h), n in imgs.most_common()))
    elif imgs:
        (w, h), n = imgs.most_common(1)[0]
        c.note(f"画像サイズ {w}x{h} ({n} 本すべて同じ)")
    if len(metas) > 1:
        c.fail("meta の image_size が揃っていない: " + " / ".join(f"{list(k)} x{n}" for k, n in metas.most_common()))
    elif metas and imgs:
        (w, h) = imgs.most_common(1)[0][0]
        mw, mh = (list(metas)[0] + (None, None))[:2]
        if (mw, mh) != (w, h):
            c.fail(f"meta の image_size {[mw, mh]} が実際の画像 {w}x{h} と違う")
    if c.status == "PASS":
        c.note(f"画像の枚数 = ステップ数 ({len(eps)} 本すべて一致)")
    return c


def check_stop(eps):
    c = Check(2, f"停止 (連続 {STOP_RUN_FAIL_SEC:.1f} 秒超で FAIL / 割合 {STOP_RATIO_WARN:.0%} 超で WARN)")
    n_stop = n_all = 0
    rows = []
    for e in eps:
        m = stop_mask(e)
        fps = fps_of(e)
        run = longest_run(m)
        ratio = float(m.mean()) if len(m) else 0.0
        n_stop += int(m.sum())
        n_all += len(m)
        rows.append((e["name"], ratio, int(m.sum()), len(m), run, run / fps))
    overall = n_stop / n_all if n_all else 0.0
    c.note(f"全体 {overall:.1%} ({n_stop}/{n_all} ステップ) / "
           f"いちばん長い連続 {max((r[5] for r in rows), default=0.0):.1f} 秒 "
           f"(上限 {STOP_RUN_FAIL_SEC:.1f} 秒 = action_horizon {ACTION_HORIZON} x {1.0 / FPS_DEFAULT:.1f} 秒)")
    if overall > STOP_RATIO_WARN:
        c.warn(f"全体の停止が {overall:.1%} で {STOP_RATIO_WARN:.0%} を超えている "
               "(停止を『動作』として覚える. build_demo_episodes.py --trim-mid-idle で詰める)")
    for name, ratio, n_s, n_t, run, sec in rows:
        msg = f"{name}: 停止 {ratio:5.1%} ({n_s:3d}/{n_t:3d})  最長 {sec:.1f} 秒 ({run} ステップ)"
        if sec > STOP_RUN_FAIL_SEC + 1e-9:
            c.fail(msg)
        elif ratio > STOP_RATIO_WARN:
            c.warn(msg)
        else:
            c.note(msg)
    return c


def check_gripper(eps):
    c = Check(3, "グリッパーの切り替え (各エピソードで「閉 → 開」の 2 回)")
    for e in eps:
        ch = grip_changes(e)
        if [w for _, w in ch] != ["閉", "開"]:
            seq = " → ".join(f"{w}(t={t})" for t, w in ch) or "切り替えなし"
            c.warn(f"{e['name']}: {len(ch)} 回  {seq}")
    if c.status == "PASS":
        c.note(f"{len(eps)} 本すべて 閉 → 開 の 2 回")
    return c


def check_limits(eps):
    c = Check(4, f"肩 j1 <= {J1_SAFE_MAX:.0f}° / j4 は J4_SAFE [{J4_SAFE[0]:+.0f}, {J4_SAFE[1]:+.0f}] の中")
    j1_top, j4_lo, j4_hi = [], [], []
    for e in eps:
        for kind in ("state", "actions"):
            v = e[kind]
            if v.shape[1] < 7:
                continue
            j1 = v[:, 1]
            i = int(j1.argmax())
            j1_top.append((float(j1[i]), e["name"], kind, i))
            if j1[i] > J1_SAFE_MAX:
                c.fail(f"{e['name']}: {kind} の j1 最大 {j1[i]:+.2f}° > {J1_SAFE_MAX:.0f}° (t={i}. "
                       "この指令はアームが黙って捨てる / LOG352)")
            j4 = v[:, 4]
            lo, hi = int(j4.argmin()), int(j4.argmax())
            j4_lo.append((float(j4[lo]), e["name"], kind, lo))
            j4_hi.append((float(j4[hi]), e["name"], kind, hi))
            if j4[lo] < J4_SAFE[0]:
                c.fail(f"{e['name']}: {kind} の j4 最小 {j4[lo]:+.2f}° < {J4_SAFE[0]:+.0f}° (t={lo})")
            if j4[hi] > J4_SAFE[1]:
                c.fail(f"{e['name']}: {kind} の j4 最大 {j4[hi]:+.2f}° > {J4_SAFE[1]:+.0f}° (t={hi})")
    if j1_top:
        v, name, kind, i = max(j1_top)
        c.note(f"j1 の最大 {v:+.2f}° ({name} {kind} t={i})")
    if j4_lo and j4_hi:
        lo, lname, lkind, li = min(j4_lo)
        hi, hname, hkind, hi_t = max(j4_hi)
        c.note(f"j4 の範囲 {lo:+.2f}° ({lname} {lkind} t={li}) .. {hi:+.2f}° ({hname} {hkind} t={hi_t})")
    return c


def check_positions(eps):
    c = Check(5, f"ブロックの位置のばらつき (つかんだ瞬間の j0 を {J0_BIN_DEG:.0f}° 刻み)")
    bins = Counter()
    for e in eps:
        close = [t for t, w in grip_changes(e) if w == "閉"]
        if not close:
            bins["不明"] += 1
            continue
        bins[round(float(e["state"][close[0], 0]) / J0_BIN_DEG) * J0_BIN_DEG] += 1
    total = sum(bins.values())
    for key in sorted(bins, key=lambda k: (isinstance(k, str), k)):
        n = bins[key]
        label = "つかんだ瞬間が不明" if isinstance(key, str) else f"j0 {key:+6.0f}° 付近"
        c.note(f"{label}: {n:2d} 本  {n / total:5.1%}  {'#' * n}")
    top = [(n, k) for k, n in bins.items() if not isinstance(k, str)]
    if top:
        n, key = max(top)
        if n / total > POS_SHARE_WARN:
            c.warn(f"j0 {key:+.0f}° 付近が {n}/{total} ({n / total:.1%}) で半分を超えている "
                   "(同じ場所ばかりだと画像を見ずに手を伸ばす / LOG348)")
    if bins.get("不明"):
        c.warn(f"つかんだ瞬間が分からないエピソードが {bins['不明']} 本 (項目 3 も見る)")
    return c


def check_action_vs_next_state(eps):
    c = Check(6, f"|action[t] - state[t+1]| の最大 (関節 6 列. {ACT_STATE_WARN_DEG}° 超で WARN)")
    worst = []
    for e in eps:
        s, a = e["state"], e["actions"]
        n = min(len(s) - 1, len(a) - 1)
        if n <= 0 or s.shape[1] < 7 or a.shape[1] < 7:
            continue
        d = np.abs(a[:n, :6] - s[1:n + 1, :6])
        t, k = np.unravel_index(int(d.argmax()), d.shape)
        worst.append((float(d[t, k]), e["name"], int(t), int(k)))
        if d[t, k] > ACT_STATE_WARN_DEG:
            c.warn(f"{e['name']}: {d[t, k]:.2f}° (t={t}, j{k}) > {ACT_STATE_WARN_DEG}° "
                   "(action は 0.1 秒後の state のはず. 停止詰めのつなぎ目を確認)")
    if worst:
        v, name, t, k = max(worst)
        c.note(f"全体の最大 {v:.2f}° ({name} t={t} j{k})")
    return c


def main():
    ap = argparse.ArgumentParser(description="π0.5 の学習前に raw_episodes を点検する (読むだけ)")
    ap.add_argument("raw_dir", help="raw_episodes のフォルダ")
    args = ap.parse_args()

    raw_dir = Path(args.raw_dir)
    if not raw_dir.is_dir():
        sys.exit(f"[ERROR] フォルダがない (KIOXIA 未接続 / パスの打ち間違い): {raw_dir}")
    eps = load_episodes(raw_dir)
    steps = sum(len(e["state"]) for e in eps)
    fps_set = sorted({fps_of(e) for e in eps})
    print(f"[DATA  ] {raw_dir}")
    print(f"         {len(eps)} 本 / {steps} ステップ / "
          f"{steps / fps_set[0]:.1f} 秒相当 / fps {'・'.join(f'{f:g}' for f in fps_set)}")
    print()

    checks = [check_shape(eps), check_stop(eps), check_gripper(eps),
              check_limits(eps), check_positions(eps), check_action_vs_next_state(eps)]
    for c in checks:
        c.report()
        print()

    n_fail = sum(1 for c in checks if c.status == "FAIL")
    n_warn = sum(1 for c in checks if c.status == "WARN")
    print("[RESULT] " + "  ".join(f"{c.num}.{c.status}" for c in checks))
    if n_fail:
        print(f"[RESULT] FAIL {n_fail} 項目 / WARN {n_warn} 項目 — このまま学習に出さない")
    elif n_warn:
        print(f"[RESULT] FAIL なし / WARN {n_warn} 項目 — 中身を見て判断する")
    else:
        print("[RESULT] 全項目 PASS")
    return 1 if n_fail else 0


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