"""OpenArm SciEdu 나뭇잎 카드 판정기: 복엽(compound) / 단엽(simple) / 없음(none: 빈 줄이거나 카드가 가려짐).
DINOv2-base 특징 + 3부류 소프트맥스 회귀. CPU로 돈다.

    from leaf_kind import LeafKind
    clf = LeafKind()                                  # classifier.json과 facebook/dinov2-base를 받는다
    r = clf.judge([frame1, frame2, frame3])           # 같은 때 찍은 d455f 640x480 RGB uint8 프레임 몇 장
    # r["kind"]: "compound" | "simple" → 옮긴다
    #            "none" → 카드가 없거나 가려졌다. 판정하지 않고 1초 뒤 다시 본다
    #            None   → 카드는 있는데 확신이 낮다(프레임 답이 갈리거나 평균 확신 < 문턱). 선생님에게 묻는다

- 사진은 카드가 줄 위에 혼자 서 있고 두 팔이 쉬고 있을 때 찍는다. 학습한 카드 사진은 카드를 옮긴 팔이 출발하기 5프레임 전이다.
- 한 프레임이라도 none이면 kind는 "none"이다. 없는 카드를 집으러 가지 않는 쪽으로 정했다.
H200 세션(leap-06)이 2026-10-04에 만들었다. 검증과 학습 자료는 README.md에 있다.
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch
from transformers import AutoImageProcessor, AutoModel

HERE = Path(__file__).resolve().parent


class LeafKind:
    def __init__(self, classifier: str | Path = HERE / "classifier.json", device: str = "cpu", threshold: float = 0.9):
        cfg = json.loads(Path(classifier).read_text(encoding="utf-8"))
        self.classes = cfg["classes"]                       # ["compound", "simple", "none"]
        self.crop = cfg["crop_xyxy"]
        self.mean, self.std = np.asarray(cfg["mean"]), np.asarray(cfg["std"])
        self.W, self.b = np.asarray(cfg["W"]), np.asarray(cfg["b"])
        self.threshold = threshold
        self.device = device
        self.processor = AutoImageProcessor.from_pretrained(cfg["feature_model"])
        self.model = AutoModel.from_pretrained(cfg["feature_model"]).to(device).eval()

    def features(self, frames: list[np.ndarray]) -> np.ndarray:
        """d455f 프레임(H x W x 3, RGB uint8, 640x480)들을 잘라 DINOv2 CLS(pooler_output)를 낸다."""
        x0, y0, x1, y1 = self.crop
        crops = [np.ascontiguousarray(f[y0:y1, x0:x1]) for f in frames]
        with torch.no_grad():
            inputs = self.processor(images=crops, return_tensors="pt").to(self.device)
            return self.model(**inputs).pooler_output.float().cpu().numpy()

    def proba(self, frames: list[np.ndarray]) -> np.ndarray:
        """프레임마다 [compound, simple, none] 확률."""
        z = ((self.features(frames) - self.mean) / self.std) @ self.W.T + self.b
        z = z - z.max(axis=1, keepdims=True)
        e = np.exp(z)
        return e / e.sum(axis=1, keepdims=True)

    def judge(self, frames: list[np.ndarray]) -> dict:
        p = self.proba(frames)
        votes = [self.classes[i] for i in p.argmax(axis=1)]
        mean = p.mean(axis=0)
        out = {"probs": dict(zip(self.classes, mean.round(4).tolist())), "votes": votes}
        if "none" in votes:
            return {"kind": "none", "confident": True, "p_compound": None, **out}
        # 카드가 있다: 복엽/단엽 둘 사이에서만 다시 나눈다
        pc = float((p[:, 0] / (p[:, 0] + p[:, 1])).mean())
        agree = len(set(votes)) == 1
        confident = agree and max(pc, 1 - pc) >= self.threshold
        kind = ("compound" if pc > 0.5 else "simple") if confident else None
        return {"kind": kind, "confident": confident, "p_compound": pc, **out}