"""DINOv2-base 특징(pooler_output)을 ONNX로 내보내고, 로컬 나뭇잎 판정기의 두 길(torch·onnxruntime)을 녹화 사진으로
견준다(이슈 #28). 로봇에 붙지 않고 GPU도 쓰지 않는다(CPU로 내보내고 잰다).

    # 1. 내보내기: torch·transformers에 onnx·onnxscript·onnxruntime을 얹은 파이썬
    CUDA_VISIBLE_DEVICES= HF_HUB_OFFLINE=1 <python> scripts/diagnostics/leaf_kind_onnx.py export --out outputs/leaf_onnx/model.onnx
    # 2. 견주기: 위 파이썬 + 녹화 leapshared/Leaf_rules_20261002_225313_subtask의 로컬 사본
    CUDA_VISIBLE_DEVICES= HF_HUB_OFFLINE=1 <python> scripts/diagnostics/leaf_kind_onnx.py compare --onnx outputs/leaf_onnx/model.onnx

저장소 venv에는 onnx·onnxruntime이 없다. 2026-10-04에는 venv의 site-packages를 .pth로 읽는 임시 venv에
onnx 1.23.1·onnxscript 0.7.2·onnx-ir 1.0.0·ml-dtypes 0.6.0·onnxruntime 1.30.0·flatbuffers 25.12.19만 더 깔아 돌렸다.

export는 facebook/dinov2-base(고정 커밋)의 pooler_output만 내는 그래프를 만든다.
- 입력은 pixel_values [N, 3, 224, 224] float32(전처리를 마친 값), 출력은 pooler_output [N, 768]이다.
- 224x224의 위치 임베딩(37x37 → 16x16 bicubic)은 torch로 한 번 셈해 상수로 넣는다. 대화층은 늘 224x224를 넣으므로
  그래프에 보간이 남지 않는다.
- 내보낸 뒤 같은 입력에서 torch와 견주고, 메타데이터(원본 모델·커밋·도구 판)를 적는다.

compare는 녹화에서 서빙 카메라처럼(RGB → BGR → JPEG 품질 80) 사진을 만들어 두 길에 넣는다.
- 검증 에피소드 6·7·46·47의 카드 22장: 결정 프레임(팔이 출발하기 5프레임 전)과 그 −12·−6 프레임으로 세 장 한 번 판정
- 같은 카드를 그리퍼가 쥐는 프레임 22장, 그 에피소드의 첫 프레임(빈 줄) 4장: 한 장씩. none이어야 한다
- 나머지 에피소드의 카드 결정 프레임과 쥐는 프레임: 한 장씩. 특징을 넓게 견준다
- 찍는 것: 전처리(numpy와 transformers의 원소 차이), 특징 코사인, 판정 일치, 판정기 만들기, 3장 판정 지연(JPEG 풀기 포함)

허브 leapshared/dinov2-base-onnx에 올린 사본은 export만 쓴다(compare는 이 저장소의 dialogue/와 녹화 도구가 있어야 한다).
"""

from __future__ import annotations

import argparse
import hashlib
import json
import statistics
import sys
import time
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve()
REPO = next((p for p in HERE.parents if (p / "pyproject.toml").is_file()), HERE.parent)  # 허브 사본은 저장소 밖
sys.path.insert(0, str(REPO))

SOURCE = "facebook/dinov2-base"
REVISION = "f9e44c814b77203eaa57a6bdbbd535f21ede1415"  # 2026-10-04 main(이 PC HF 캐시)
SIZE = 224
DATASET = "Leaf_rules_20261002_225313_subtask"
HELD_OUT = (6, 7, 46, 47)
KIND = {"복엽": "compound", "단엽": "simple"}


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def cosines(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    a, b = np.asarray(a, np.float64), np.asarray(b, np.float64)
    return (a * b).sum(1) / (np.linalg.norm(a, axis=1) * np.linalg.norm(b, axis=1))


# ── export ──
def export(out: Path, revision: str, opset: int | None) -> None:
    import onnx
    import onnxruntime as ort
    import torch
    import transformers
    from transformers import AutoModel

    torch.manual_seed(0)
    model = AutoModel.from_pretrained(SOURCE, revision=revision).eval()
    probe = torch.randn(3, 3, SIZE, SIZE)
    with torch.no_grad():
        want = model(pixel_values=probe).pooler_output  # 고치기 전 모델의 답

    class Pooler(torch.nn.Module):
        """pixel_values [N, 3, 224, 224] → pooler_output [N, 768]. 위치 임베딩은 224x224로 미리 셈한 상수다."""

        def __init__(self, inner) -> None:
            super().__init__()
            self.inner = inner
            embeddings = inner.embeddings
            tokens = torch.zeros(1, 1 + (SIZE // embeddings.patch_size) ** 2, inner.config.hidden_size)
            with torch.no_grad():
                self.register_buffer("pos", embeddings.interpolate_pos_encoding(tokens, SIZE, SIZE).clone())
            embeddings.interpolate_pos_encoding = lambda _embeddings, _height, _width: self.pos

        def forward(self, pixel_values):
            return self.inner(pixel_values=pixel_values).pooler_output

    pooler = Pooler(model).eval()
    with torch.no_grad():
        got = pooler(probe)
    print(f"상수 위치 임베딩: 원래 모델과 최대 차이 {float((got - want).abs().max()):.3g}")

    out.parent.mkdir(parents=True, exist_ok=True)
    batch = torch.export.Dim("batch", min=1, max=64)
    torch.onnx.export(pooler, (probe,), str(out), input_names=["pixel_values"], output_names=["pooler_output"],
                      dynamic_shapes={"pixel_values": {0: batch}}, opset_version=opset, dynamo=True,
                      external_data=False)
    graph = onnx.load(str(out))
    for node in graph.graph.node:  # 노드마다 붙는 stack_trace 등: 내보낸 PC의 경로·줄 번호라 계산과 상관없다. 지우면 파일이 늘 같다
        del node.metadata_props[:]
    onnx.helper.set_model_props(graph, {
        "source_model": SOURCE, "source_revision": revision, "output": "pooler_output",
        "input": f"pixel_values [N, 3, {SIZE}, {SIZE}] float32: AutoImageProcessor 전처리를 마친 값",
        "position_embeddings": f"{SIZE}x{SIZE} 입력으로 미리 셈한 상수",
        "exported_with": f"torch {torch.__version__}, transformers {transformers.__version__}, onnx {onnx.__version__}",
        "made_by": "LEAP-LAB-KUS/openarmsciedu-robot scripts/diagnostics/leaf_kind_onnx.py export",
    })
    onnx.save(graph, str(out))
    onnx.checker.check_model(str(out))
    opsets = {o.domain or "ai.onnx": o.version for o in graph.opset_import}
    session = ort.InferenceSession(str(out), providers=["CPUExecutionProvider"])
    for n in (1, 3):
        ort_out = session.run(None, {"pixel_values": probe[:n].numpy()})[0]
        diff = np.abs(ort_out - want[:n].numpy()).max()
        print(f"onnxruntime {ort.__version__}, {n}장: 최대 차이 {diff:.3g}, 코사인 최소 "
              f"{cosines(ort_out, want[:n].numpy()).min():.8f}")
    print(f"썼다: {out} ({out.stat().st_size / 1e6:.1f} MB, sha256 {sha256(out)}, opset {opsets})")


# ── compare ──
def serving_jpeg(image) -> bytes:
    """서빙 카메라(serving/cameras.py)처럼: RGB → BGR → JPEG 품질 80."""
    import cv2

    rgb = np.asarray(image.convert("RGB"), dtype=np.uint8)
    ok, buf = cv2.imencode(".jpg", cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_JPEG_QUALITY, 80])
    assert ok
    return buf.tobytes()


def photos(root: Path) -> list[dict]:
    """녹화의 사진들: {"name", "kind", "group", "held_out", "jpeg"}. 검증 카드는 세 장(−12·−6·0)이 한 판정이다."""
    sys.path.insert(0, str(REPO / "scripts" / "diagnostics"))
    sys.path.insert(0, str(REPO / "scripts" / "verifier"))
    import common
    from leaf_view import all_cards

    meta = common.episodes_meta(root)
    by_episode: dict[int, list[dict]] = {}
    for card in all_cards(root):
        by_episode.setdefault(card["episode"], []).append(card)
    rows = []
    for episode, cards in sorted(by_episode.items()):
        held = episode in HELD_OUT
        want = {0} if held else set()
        for c in cards:
            want |= {c["decide"], c["close"]} | ({max(0, c["decide"] - 12), max(0, c["decide"] - 6)} if held else set())
        got = common.read_frames(root, meta, episode, common.CHEST, sorted(want))
        if held:
            rows.append({"name": f"ep{episode:02d} 첫 프레임", "kind": "none", "group": "first", "held_out": True,
                         "jpeg": [serving_jpeg(got[0])]})
        for c in cards:
            tag = f"ep{episode:02d}#{c['card']}"
            offsets = (-12, -6, 0) if held else (0,)
            rows.append({"name": tag, "kind": KIND[c["kind"]], "group": "card", "held_out": held,
                         "jpeg": [serving_jpeg(got[max(0, c["decide"] + o)]) for o in offsets]})
            rows.append({"name": f"{tag} 쥠", "kind": "none", "group": "close", "held_out": held,
                         "jpeg": [serving_jpeg(got[c["close"]])]})
    return rows


def compare(onnx_path: Path, root: Path, threads: int, repeat: int, out: Path | None) -> None:
    import torch
    import transformers

    from dialogue.leaf_kind.leaf_kind import LeafKind
    from dialogue.leaf_local import LocalLeafJudge, rgb_frame
    from dialogue.leaf_onnx import OnnxLeafKind, pixel_values

    classifier = REPO / "dialogue" / "leaf_kind" / "classifier.json"
    torch.set_num_threads(threads)
    started = time.perf_counter()
    reference = LeafKind(classifier=classifier, device="cpu", threshold=0.9)
    torch_load = time.perf_counter() - started
    started = time.perf_counter()
    candidate = OnnxLeafKind(onnx_path, classifier, threshold=0.9, threads=threads)
    onnx_load = time.perf_counter() - started

    rows = photos(root)
    x0, y0, x1, y1 = reference.crop
    pixel_diff = pixel_total = 0
    feats_t, feats_o, frame_rows, judged = [], [], [], []
    for row in rows:
        frames = [rgb_frame(j) for j in row["jpeg"]]
        for frame in frames:  # 한 장씩: 전처리와 특징
            crop = np.ascontiguousarray(frame[y0:y1, x0:x1])
            ref_px = reference.processor(images=[crop], return_tensors="np")["pixel_values"]
            pixel_diff += int((ref_px != pixel_values([crop])).sum())
            pixel_total += ref_px.size
            ft, fo = reference.features([frame]), candidate.features([frame])
            feats_t.append(ft[0])
            feats_o.append(fo[0])
            pt, po = reference.proba([frame])[0], candidate.proba([frame])[0]
            frame_rows.append({"name": row["name"], "kind": row["kind"], "torch": reference.classes[int(pt.argmax())],
                               "onnx": candidate.classes[int(po.argmax())], "dp": float(np.abs(pt - po).max())})
        a, b = reference.judge(frames), candidate.judge(frames)  # 대화층처럼 한 번에
        judged.append({"name": row["name"], "kind": row["kind"], "group": row["group"], "held_out": row["held_out"],
                       "frames": len(frames), "torch": a, "onnx": b})

    cos = cosines(np.stack(feats_t), np.stack(feats_o))
    absdiff = np.abs(np.stack(feats_t) - np.stack(feats_o))
    same_vote = sum(r["torch"] == r["onnx"] for r in frame_rows)
    same_judge = sum((j["torch"]["kind"], j["torch"]["votes"]) == (j["onnx"]["kind"], j["onnx"]["votes"]) for j in judged)
    print(f"판정기 만들기(꾸러미를 불러온 뒤): torch {torch_load:.2f} s, onnx {onnx_load:.2f} s (스레드 {threads})")
    print(f"전처리: 다른 원소 {pixel_diff}/{pixel_total}")
    print(f"특징 {len(cos)}장: 코사인 최소 {cos.min():.12f}(1 − 코사인 최대 {1 - cos.min():.2g}), 원소 최대 차이 "
          f"{absdiff.max():.3g}(특징 원소 크기 중앙 {np.median(np.abs(np.stack(feats_t))):.3g}), 세 부류 확률 최대 차이 "
          f"{max(r['dp'] for r in frame_rows):.3g}")
    print(f"한 장 판정(argmax) 같음 {same_vote}/{len(frame_rows)}, 판정(kind·votes) 같음 {same_judge}/{len(judged)}")

    def tally(select, label):
        chosen = [j for j in judged if select(j)]
        for path in ("torch", "onnx"):
            right = sum(j[path]["kind"] == j["kind"] for j in chosen)
            other = {}
            for j in chosen:
                if j[path]["kind"] != j["kind"]:
                    other[str(j[path]["kind"])] = other.get(str(j[path]["kind"]), 0) + 1
            print(f"  {label} {path}: 맞음 {right}/{len(chosen)}" + (f", 다른 답 {other}" if other else ""))

    tally(lambda j: j["held_out"] and j["group"] == "card", "검증 카드(세 장 한 번)")
    tally(lambda j: j["held_out"] and j["group"] == "close", "검증 쥐는 프레임")
    tally(lambda j: j["held_out"] and j["group"] == "first", "검증 첫 프레임(빈 줄)")
    tally(lambda j: not j["held_out"] and j["group"] == "card", "나머지 카드(한 장)")
    tally(lambda j: not j["held_out"] and j["group"] == "close", "나머지 쥐는 프레임")

    # 지연: 대화층과 같은 길(JPEG 세 장 → 풀기 → 자르기 → 전처리 → 특징 → 판정)
    three = next(r for r in rows if r["held_out"] and r["group"] == "card")["jpeg"]
    latency = {}
    for path, clf in (("torch", reference), ("onnx", candidate)):
        judge = LocalLeafJudge(clf, path)
        for _ in range(3):
            judge.judge(three)
        times = []
        for _ in range(repeat):
            started = time.perf_counter()
            judge.judge(three)
            times.append((time.perf_counter() - started) * 1000)
        latency[path] = times
        print(f"3장 판정 {path}: 중앙 {statistics.median(times):.0f} ms, 최대 {max(times):.0f} ms ({repeat}번)")
    if out:
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(json.dumps({
            "onnx": str(onnx_path), "onnx_sha256": sha256(onnx_path), "threads": threads,
            "versions": {"torch": torch.__version__, "transformers": transformers.__version__,
                         "onnxruntime": __import__("onnxruntime").__version__},
            "load_s": {"torch": torch_load, "onnx": onnx_load}, "pixel_diff": [pixel_diff, pixel_total],
            "cosine": {"min": float(cos.min()), "median": float(np.median(cos))}, "feature_absdiff_max": float(absdiff.max()),
            "frames": frame_rows, "judged": judged, "latency_ms": latency}, ensure_ascii=False, indent=1))
        print("썼다:", out)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="command", required=True)
    ex = sub.add_parser("export", help="dinov2-base pooler_output → ONNX(CPU)")
    ex.add_argument("--out", type=Path, default=REPO / "outputs" / "leaf_onnx" / "model.onnx")
    ex.add_argument("--revision", default=REVISION)
    ex.add_argument("--opset", type=int, default=None, help="비우면 torch 내보내기의 기본값")
    co = sub.add_parser("compare", help="녹화 사진으로 torch 길과 onnxruntime 길 견주기")
    co.add_argument("--onnx", type=Path, required=True)
    co.add_argument("--root", type=Path, default=Path.home() / ".cache" / "huggingface" / "lerobot" / "leapshared" / DATASET)
    co.add_argument("--threads", type=int, default=4, help="대화층 설정 leaf.local.threads와 같게")
    co.add_argument("--repeat", type=int, default=30)
    co.add_argument("--out", type=Path, default=REPO / "outputs" / "leaf_onnx" / "compare.json")
    args = parser.parse_args()
    if args.command == "export":
        export(args.out, args.revision, args.opset)
    else:
        compare(args.onnx, args.root, args.threads, args.repeat, args.out)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())