#!/usr/bin/env python3
"""count_by_position.py — 学習データの本数を, ブロックの置き場所 (LOG352 の教えた位置 1〜5) ごとに数える (2026-11-07 依頼 118).

エイイチの決め (2026-11-07): 置き場所ごとの本数をそろえる. 目安は各位置 10 本.
撮り足す本数を決めるときと, merge_raw_episodes.py でデータセットを作り終えたときの README の表に使う.

入力 (どちらでも. いくつでも):
  データセット名            datasets/{名前}/raw_episodes を数える
  raw_episodes のフォルダ   そのフォルダを数える (ひとつ上の index.json があれば外した本も数える)

位置の決め方:
  教えた動き (run_pick_place の run):
    試行フォルダの teach/ にある教えた記録のファイル名で決める (LOG352 の README の表と同じ対応. TEACH_FILES).
    teach/ がない・表にないファイル名なら「分からない」.
  補正入り (pi05_policy の run. meta.json の source が pi05_policy で始まる):
    上カメラ (recording/top.mp4) の最初のコマでマットの四隅を見つけ, マットの上の mm に直す対応を作る
    (依頼 107 twin_prep と同じ. visual_servo_pick の四隅 (±310, −89)・(±310, 361) mm).
    爪を閉じる命令 (いちばん初め) より前のコマを 3 つおきに見て, ブロックらしい赤 (twin_prep.red_candidates と同じ.
    根元から 120mm 以内・細長い・埋まりの少ない赤は外す) がいちばん大きく写ったコマの中心を, ブロックの位置とする
    (HOME の最初のコマでは爪の下に半分隠れていることがあるため).
    その位置から TEACH_REF_MM のいちばん近い位置を選び, 離れが MAX_DIST_MM (30mm) を超えたら「分からない」.
    四隅・ブロックが見つからないときも「分からない」(理由を残す).
  TEACH_REF_MM: 同じ見つけ方を LOG352 で採った 19 本 (10/19) にかけた値の位置ごとの中央値.
    位置 1 は 183454 の 1 本だけ (爪の陰で赤が小さく写る. 866px). 位置 2 は 191257 が見つからず 4 本,
    位置 3 は 192539 が見つからず 3 本. 位置 4 から 46mm ずらして置いた 193729 は (237.4, 52.6) で, 位置 4 から 48mm
    (使わない試行. この線なら「分からない」になる).
    教えた位置どうしのいちばん近い間は位置 1〜2 の 63mm. 線 30mm はその半分より少し下.
    座標はマットの上の mm (visual_servo_pick の向き). 依頼 107 C の Unity の値 (T1〜T5) とは原点・向きが違う.

外した本・手の本 (データセットの index.json があるとき):
  merge_raw_episodes.py の index.json の excluded / build_*_episodes.py の index.json の used=false の本を,
  同じ決め方で位置に割り当て, 「浅くて外した」(grasp_depth.shallow か理由が「浅い」「shallow」で始まる) と
  「手で外した」(hand_in_frame か grasp_depth.hand_suspect) を数える.
  データに入っている本の「手が写っていそう」は index に値があるときだけ数える (--measure-depth を付けると
  grasp_depth.py で 1 本ずつ測る. 遅い).

  python3 pi05/count_by_position.py mecharm_redblock_box_teach_v1_trim2_corr_v5_deep
  python3 pi05/count_by_position.py A B --target 10 --json out.json

読むだけ (記録・データセットには何も書かない. --json のファイルだけ書く).
"""
import argparse
import contextlib
import json
import math
import sys
from pathlib import Path

import cv2
import numpy as np

_HERE = Path(__file__).resolve()
sys.path.insert(0, str(_HERE.parent))
sys.path.insert(0, str(_HERE.parent.parent))   # mat_corner_detection.py (r2r2r-scripts の直下)

import grasp_depth as gd  # noqa: E402

DATA_ROOT = Path("/Volumes/KIOXIA/r2r2r-data")
RUN_SUBDIRS = ("run_pick_place", "pi05_policy")
POSITIONS = (1, 2, 3, 4, 5)
UNKNOWN = "unknown"
DEFAULT_TARGET = 10

# LOG352 の README の表 (runs/experiments/record_teach_demo/20261019_LOG352/README.md)
TEACH_FILES = {
    "drag_v2_20221013_144348.json": 1,
    "drag_v2_20261109_094222.json": 1,     # 位置 1 を教え直した記録 (DEEP_DEMOS 2-1・LOG463)
    "drag_v2_20221013_230205.json": 2,
    "drag_v2_20221013_231249.json": 3,
    "drag_v2_20261109_123000.json": 3,     # 位置 3 を教え直した記録 (DEEP_DEMOS 2-1・2026-11-09)
    "drag_v2_20221013_232345.json": 4,
    "drag_v2_20261109_130000.json": 4,     # 位置 4 のブロックの上を差し替えた記録 (2026-11-09)
    "drag_v2_20261109_130300.json": 4,     # 位置 4 をゲームパッドで教え直した記録 (DEEP_DEMOS 2-1・2026-11-09)
    "drag_v2_20221013_233323.json": 5,
}
# 教えた位置のブロックの中心 (マットの上の mm). 決め方は頭の説明
TEACH_REF_MM = {
    1: (182.7, 188.8),
    2: (215.5, 242.2),
    3: (123.6, 235.6),
    4: (240.8, 101.3),
    5: (-48.8, 246.4),
}
MAX_DIST_MM = 30.0
MAT_CORNERS_MM = [[310, -89], [-310, -89], [-310, 361], [310, 361]]   # visual_servo_pick と同じ (twin_prep)
BLOCK_MIN_DIST_FROM_BASE_MM = 120.0                                   # record_teach_demo と同じ (twin_prep)
FRAME_STEP = 3


# ---------------------------------------------------------------- 位置を決める
def teach_position(run_dir):
    """教えた動きの run: teach/ の教えた記録のファイル名から. 戻り値 (位置 or UNKNOWN, 根拠の文字)."""
    tdir = Path(run_dir) / "teach"
    files = sorted(p.name for p in tdir.glob("[!.]*.json")) if tdir.is_dir() else []
    if not files:
        return UNKNOWN, "teach/ に教えた記録がない"
    pos = {TEACH_FILES[f] for f in files if f in TEACH_FILES}
    if len(pos) != 1:
        return UNKNOWN, f"teach/ のファイルが表と合わない: {', '.join(files)}"
    return pos.pop(), f"teach/{files[0]}"


def nearest_position(mm, max_dist=MAX_DIST_MM):
    """ブロックの位置 (mm) → (位置 or UNKNOWN, いちばん近い位置, 離れ mm)."""
    p, d = min(((k, math.hypot(mm[0] - v[0], mm[1] - v[1])) for k, v in TEACH_REF_MM.items()), key=lambda x: x[1])
    return (p if d <= max_dist else UNKNOWN), p, round(d, 1)


def _px_to_mm(H, x, y):
    r = cv2.perspectiveTransform(np.array([[[float(x), float(y)]]], np.float32), H)
    return float(r[0][0][0]), float(r[0][0][1])


def red_candidates(frame, mat_bbox, H):
    """twin_prep.red_candidates と同じ (向きは使わないので出さない)."""
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    m = cv2.inRange(hsv, np.array([0, 100, 100]), np.array([10, 255, 255])) | \
        cv2.inRange(hsv, np.array([160, 100, 100]), np.array([180, 255, 255]))
    m = cv2.morphologyEx(m, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8))
    cnts, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    out = []
    for c in cnts:
        a = cv2.contourArea(c)
        if a < 500:
            continue
        (cx, cy), (rw, rh), _ang = cv2.minAreaRect(c)
        if not (mat_bbox[0] <= cx <= mat_bbox[1] and mat_bbox[2] <= cy <= mat_bbox[3]):
            continue
        aspect = max(rw, rh) / max(1e-6, min(rw, rh))
        extent = a / max(1e-6, rw * rh)
        mm = _px_to_mm(H, cx, cy)
        if math.hypot(*mm) < BLOCK_MIN_DIST_FROM_BASE_MM or aspect > 1.8 or extent < 0.6:
            continue
        out.append({"area": float(a), "px": (float(cx), float(cy)), "mm": mm})
    return out


def detect_corners(frame):
    with contextlib.redirect_stdout(sys.stderr):
        import mat_corner_detection as mcd
        return mcd.detect_corners(frame)


def block_mm(run_dir, corners_fn=None):
    """補正入りの run: 上カメラで見つけたブロックの位置. 戻り値 (dict or None, 理由)."""
    corners_fn = corners_fn or detect_corners
    rec = Path(run_dir) / "recording"
    if not (rec / "top.mp4").exists() or not (rec / "top_frame_log.csv").exists():
        return None, "top.mp4 か top_frame_log.csv がない"
    tc = None
    if (rec / "states.csv").exists():
        gt, gs = gd.load_gripper_cmds(rec / "states.csv")
        tc = gd.close_time(gt, gs)
    _, fts = gd.load_frame_log(rec / "top_frame_log.csv")
    cap = cv2.VideoCapture(str(rec / "top.mp4"))
    try:
        ok, f0 = cap.read()
        if not ok:
            return None, "top.mp4 が読めない"
        corners = corners_fn(f0)
        if corners is None:
            return None, "マットの四隅が見つからない"
        co = [[float(x), float(y)] for x, y in corners]
        H = cv2.getPerspectiveTransform(np.array(co, np.float32), np.array(MAT_CORNERS_MM, np.float32))
        bbox = (int(min(co[0][0], co[3][0])), int(max(co[1][0], co[2][0])),
                int(min(co[0][1], co[1][1])), int(max(co[2][1], co[3][1])))
        best, k, f = None, 0, f0
        while ok:
            if k < len(fts) and tc is not None and fts[k] > tc:
                break
            if k % FRAME_STEP == 0:
                cs = red_candidates(f, bbox, H)
                if cs:
                    c = max(cs, key=lambda c: c["area"])
                    if best is None or c["area"] > best["area"]:
                        best = dict(c, frame=k)
            ok, f = cap.read()
            k += 1
    finally:
        cap.release()
    if best is None:
        return None, "閉じる前のコマでブロックが見つからない"
    return best, "ok"


def find_run_dir(stem, runs_root):
    for sub in RUN_SUBDIRS:
        d = Path(runs_root) / sub / stem
        if d.is_dir():
            return d
    return None


def is_correction(meta, stem):
    src = (meta or {}).get("source", "")
    if src:
        return src.startswith("pi05_policy")
    return stem.startswith("pi05_")


def classify(stem, meta, runs_root, corners_fn=None):
    """1 本の位置. 戻り値 dict(stem, kind, position, how, ...)."""
    kind = "correction" if is_correction(meta, stem) else "teach"
    out = {"stem": stem, "kind": kind, "position": UNKNOWN, "how": None}
    rd = find_run_dir(stem, runs_root)
    if rd is None:
        out["how"] = "runs/ に試行フォルダがない"
        return out
    if kind == "teach":
        out["position"], out["how"] = teach_position(rd)
        return out
    try:
        b, why = block_mm(rd, corners_fn)
    except Exception as e:   # 見つけ方で落ちても数えるのは止めない
        b, why = None, f"exception:{type(e).__name__}:{e}"
    if b is None:
        out["how"] = why
        return out
    pos, near, d = nearest_position(b["mm"])
    out.update(position=pos, how=f"top.mp4 のブロック (コマ {b['frame']}) → 位置 {near} まで {d}mm",
               block_mm=[round(v, 1) for v in b["mm"]], nearest=near, dist_mm=d)
    if pos == UNKNOWN:
        out["how"] += f" (線 {MAX_DIST_MM:g}mm を超える)"
    return out


# ---------------------------------------------------------------- 外した本・手の本
def _shallow(v):
    g = v.get("grasp_depth") or {}
    r = str(v.get("reason", ""))
    return bool(g.get("shallow")) or r.startswith("浅い") or r.startswith("shallow")


def _hand(v):
    g = v.get("grasp_depth") or {}
    return bool(v.get("hand_in_frame")) or bool(g.get("hand_suspect"))


def excluded_entries(ds_dir):
    """データセットの index.json から, 外した本 {名前: 中身}. merge の excluded と build の used=false の両方."""
    p = Path(ds_dir) / "index.json"
    if not p.exists():
        return {}, {}
    idx = json.loads(p.read_text())
    out, kept_info = {}, {}
    for k, v in (idx.get("excluded") or {}).items():
        out[k] = v
    for k, v in idx.items():
        if isinstance(v, dict) and "used" in v and "run" in v:
            if v["used"]:
                kept_info[k] = v
            elif v.get("grasp_depth") or v.get("hand_in_frame"):
                out[k] = v
    return out, kept_info


# ---------------------------------------------------------------- 数える
def resolve_input(x, data_root):
    p = Path(x)
    if p.is_dir() and (p.name == "raw_episodes" or list(p.glob("[!.]*/episode.npz"))):
        return p, p.parent
    d = Path(data_root) / "datasets" / x
    return d / "raw_episodes", d


def count(inputs, data_root=DATA_ROOT, target=DEFAULT_TARGET, measure_depth=False, corners_fn=None, log=print):
    runs_root = Path(data_root) / "runs"
    eps, excl, kept_info, seen = [], [], {}, set()
    for x in inputs:
        raw, ds = resolve_input(x, data_root)
        if not raw.is_dir():
            raise SystemExit(f"[ABORT] raw_episodes がない: {raw}")
        for ep in sorted(p.parent for p in raw.glob("[!.]*/episode.npz")):
            if ep.name in seen:
                continue
            seen.add(ep.name)
            meta = json.loads((ep / "meta.json").read_text()) if (ep / "meta.json").exists() else {}
            stem = meta.get("run_stem", ep.name)
            eps.append((ep.name, stem, meta))
        ex, ki = excluded_entries(ds)
        kept_info.update(ki)
        for k, v in ex.items():
            if k not in seen:
                seen.add(k)
                excl.append((k, v))
    rows = []
    for name, stem, meta in eps:
        r = classify(stem, meta, runs_root, corners_fn)
        r["name"] = name
        g = (kept_info.get(stem) or {}).get("grasp_depth")
        if measure_depth:
            rd = find_run_dir(stem, runs_root)
            g = gd.measure_run(rd) if rd else None
        r["hand_suspect"] = bool(g and g.get("hand_suspect"))
        r["depth_known"] = g is not None
        rows.append(r)
        log(f"[POS ] {name}  {'教えた動き' if r['kind'] == 'teach' else '補正入り'}  位置 {r['position']}  {r['how']}")
    ex_rows = []
    for name, v in excl:
        kind_meta = {"source": "pi05_policy"} if (v.get("kind") == "補正入り" or name.startswith("pi05_")) else {}
        r = classify(name, kind_meta, runs_root, corners_fn)
        r.update(name=name, shallow=_shallow(v), hand=_hand(v), reason=v.get("reason"))
        ex_rows.append(r)
        log(f"[EXCL] {name}  位置 {r['position']}  {r.get('reason')}")
    return summarize(rows, ex_rows, target)


def summarize(rows, ex_rows, target):
    keys = list(POSITIONS) + [UNKNOWN]
    table = {}
    for k in keys:
        rs = [r for r in rows if r["position"] == k]
        es = [r for r in ex_rows if r["position"] == k]
        n = len(rs)
        table[k] = {
            "episodes": n,
            "teach": sum(r["kind"] == "teach" for r in rs),
            "correction": sum(r["kind"] == "correction" for r in rs),
            "to_add": (max(0, target - n) if k != UNKNOWN else None),
            "excluded_shallow": sum(r["shallow"] for r in es),
            "excluded_hand": sum(r["hand"] for r in es),
            "hand_suspect_in_data": sum(r["hand_suspect"] for r in rs),
        }
    return {"target": target, "max_dist_mm": MAX_DIST_MM, "table": table,
            "total": len(rows), "depth_known": sum(r["depth_known"] for r in rows),
            "episodes": rows, "excluded": ex_rows}


def table_lines(res):
    """Markdown の表 (README と端末で同じもの)."""
    t = res["table"]
    lines = [f"| 位置 | 本数 | 教えた動き | 補正入り | 目標 {res['target']} 本まで | 浅くて外した | 手で外した | "
             "入っている本で手が写っていそう |",
             "| --- | --- | --- | --- | --- | --- | --- | --- |"]
    for k, v in t.items():
        name = f"位置 {k}" if k != UNKNOWN else "分からない"
        add = "—" if v["to_add"] is None else f"+{v['to_add']}"
        hand_in = v["hand_suspect_in_data"] if res["depth_known"] else "—"
        lines.append(f"| {name} | {v['episodes']} | {v['teach']} | {v['correction']} | {add} | "
                     f"{v['excluded_shallow']} | {v['excluded_hand']} | {hand_in} |")
    lines.append(f"| 合計 | {res['total']} | {sum(v['teach'] for v in t.values())} | "
                 f"{sum(v['correction'] for v in t.values())} | "
                 f"+{sum(v['to_add'] or 0 for v in t.values())} | "
                 f"{sum(v['excluded_shallow'] for v in t.values())} | {sum(v['excluded_hand'] for v in t.values())} | "
                 f"{sum(v['hand_suspect_in_data'] for v in t.values()) if res['depth_known'] else '—'} |")
    return lines


def readme_section(res):
    unk = [r for r in res["episodes"] + res["excluded"] if r["position"] == UNKNOWN]
    lines = ["## 置き場所ごとの本数 (`pi05/count_by_position.py`。依頼 118)", "",
             "教えた動きは `teach/` の教えた記録のファイル名 (LOG352 の表) で、補正入りは上カメラで見つけたブロックの位置を"
             f"いちばん近い教えた位置に割り当てた ({res['max_dist_mm']:g}mm を超えたら「分からない」)。"
             "「浅くて外した」「手で外した」は、このデータセットの `index.json` で外したエピソードの数 (データには入っていない)。", ""]
    lines += table_lines(res)
    if not res["depth_known"]:
        lines += ["", "- 入っている本の「手が写っていそう」は測っていない (`--measure-depth` で測る)"]
    for r in unk:
        lines.append(f"- 分からない: {r['name']} ({r['how']})")
    return lines


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("inputs", nargs="+", help="データセット名 (datasets/ の下) か raw_episodes のフォルダ")
    ap.add_argument("--target", type=int, default=DEFAULT_TARGET, help=f"位置ごとの目標本数 (既定 {DEFAULT_TARGET})")
    ap.add_argument("--measure-depth", action="store_true",
                    help="入っている本を grasp_depth.py で測り, 手が写っていそうな本を数える (遅い)")
    ap.add_argument("--json", help="結果を JSON に書く (書く先はこのファイルだけ)")
    ap.add_argument("--data-root", default=str(DATA_ROOT), help=argparse.SUPPRESS)
    args = ap.parse_args()
    res = count(args.inputs, args.data_root, args.target, args.measure_depth)
    print()
    print("\n".join(table_lines(res)))
    if args.json:
        Path(args.json).write_text(json.dumps(res, indent=2, ensure_ascii=False, default=str))
        print(f"[JSON] {args.json}")
    return 0


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