src/poolcoach_cv/ballid.py

"""BallID — gán SỐ bi 9-ball hai tầng theo MÀU, không train (06/08/2026).

Thực thi `PoolCoach_CV_BallID_Design.md` (chốt 06/08): dataset pix2pockets
KHÔNG có nhãn số nên tầng 1 là baseline màu thuần numpy; khi có bộ ảnh VN
gán nhãn sẽ swap classifier train cùng chữ ký (crop → vector cost 9 số),
tầng 2 giữ nguyên.

- **Tầng 1** (`score_balls`): mỗi bbox bi → vector cost 9 số. Crop co 60%
  bbox quanh tâm (né nỉ + bóng) → chuẩn hoá trắng theo median crop bi CUE
  (đèn vàng quán VN ám cue bao nhiêu thì ám các bi khác bấy nhiêu; không
  cue → bỏ, cờ ``wb=False``) → tách pixel TRẮNG/MÀU → cost = khoảng cách
  hue circular (trọng số S/V thấp hơn H) tới palette + phạt lệch
  ``white_ratio``; bi 8 chấm theo V thấp thay vì hue.
- **Tầng 2** (`assign_numbers`): ma trận cost → gán greedy TẤT ĐỊNH với
  ràng buộc mỗi số ≤ 1 bi (duyệt cặp cost tăng dần; hoà cost → index nhỏ
  trước); cost > ``cost_max`` → không gán (``None``). ``number_conf`` =
  nghịch biến cost + độ tách với ứng viên nhì.

Thuần numpy như `homography.py` — KHÔNG import cv2/torch → unit test chạy
trong cả `poolcoach-env` lẫn `poolcoach-cv-env`.

**Thứ tự kênh màu là BGR (quy ước OpenCV)** — mọi ảnh/crop vào đây đều
BGR uint8 (hoặc float [0,255]). Lẫn RGB↔BGR làm bi 2 (xanh dương) thành
cam và ngược lại — test synthetic viết màu BGR tường minh để khoá đúng
điều này.
"""

from __future__ import annotations

import numpy as np

_EPS = 1e-12

# ---------------------------------------------------------------------------
# BẢNG HẰNG DUY NHẤT của BallID — palette + mọi ngưỡng (design §2).
#
# Căn cứ: màu bi 9-ball chuẩn (1 vàng · 2 xanh dương · 3 đỏ · 4 tím · 5 cam ·
# 6 xanh lá · 7 đỏ nâu · 8 đen · 9 sọc trắng-vàng), giá trị H/S/V ước lượng
# tay từ bộ bi tiêu chuẩn dưới đèn trung tính. Quy ước HSV NỘI BỘ của module:
# H ∈ [0,360)°, S/V ∈ [0,1] (KHÔNG phải hue OpenCV [0,180) — module không
# dùng cv2). KHÔNG chỉnh các số này sau khi đã thấy số đo ảnh thật: thấy số
# rồi muốn chỉnh là vòng 2, ghi HANDOFF để Cowork quyết (design §4).
BALLID = {
    # Palette màu đại diện pixel MÀU (median HSV của crop): số → (h, s, v).
    # Bi 8 KHÔNG có mục ở đây — đen không có hue tin được, chấm riêng theo
    # V thấp bằng bộ ngưỡng v8/s8 bên dưới. Bi 9 cùng hue vàng với bi 1 —
    # phân biệt bằng white_ratio (sọc trắng), không phải bằng màu.
    "palette": {
        1: (52.0, 0.85, 0.85),   # vàng
        2: (225.0, 0.80, 0.55),  # xanh dương
        3: (5.0, 0.90, 0.70),    # đỏ
        4: (285.0, 0.65, 0.45),  # tím
        5: (27.0, 0.90, 0.80),   # cam
        6: (135.0, 0.75, 0.45),  # xanh lá
        7: (350.0, 0.75, 0.40),  # đỏ nâu (hue sát bi 3, tối + xỉn hơn)
        9: (52.0, 0.85, 0.85),   # phần vàng của sọc
    },
    # Khoảng white_ratio kỳ vọng [lo, hi] — lệch ra ngoài phạt tuyến tính
    # theo khoảng cách tới biên (design §2: trơn ~0–0.15, sọc 9 ~0.3–0.6).
    "white_range": {"solid": (0.0, 0.15), "eight": (0.0, 0.10),
                    "stripe": (0.30, 0.60)},
    # Trọng số cost: hue nặng nhất, S/V phụ (design §2); white ngang hue để
    # cặp 1↔9 (cùng hue) tách được bằng sọc.
    "w_hue": 1.0, "w_sat": 0.25, "w_val": 0.25, "w_white": 1.0,
    # Bi 8: cost = w8_val·max(0, v−v8_max) + w8_sat·max(0, s−s8_max) + phạt
    # trắng. Bi đen chuẩn v≈0.1–0.3; hệ số dốc để bi màu (v≈0.5+) không với
    # tới số 8 dưới ngưỡng cost_max.
    "v8_max": 0.35, "s8_max": 0.40, "w8_val": 3.0, "w8_sat": 1.0,
    # Pixel TRẮNG = sáng cao + bão hoà thấp (sọc/bóng loé): V ≥ white_v_min
    # VÀ S ≤ white_s_max. Sau WB, trắng thật S ≈ 0.0–0.15; ám vàng KHÔNG WB
    # đẩy S lên ~0.5 → trắng "biến mất" — chính là cái đối chứng âm đo.
    "white_v_min": 0.65, "white_s_max": 0.25,
    # Crop co 60% kích thước bbox quanh tâm — né nỉ bàn + bóng đổ (design §2).
    "crop_shrink": 0.60,
    # Dưới tỉ lệ pixel MÀU này thì median màu hết nghĩa (bi gần toàn trắng —
    # vd bi trắng thứ hai bị hạ từ cue xuống ball) → dùng median MỌI pixel;
    # white_ratio cao sẽ tự đẩy cost vượt cost_max → number=None.
    "min_color_frac": 0.05,
    # Tầng 2: cost gán > cost_max → không gán (None); margin_ref chuẩn hoá
    # độ tách với ứng viên nhì khi tính conf.
    "cost_max": 0.60, "margin_ref": 0.15,
    # number_conf = w_conf_cost·(1 − cost/cost_max)
    #             + w_conf_margin·min(1, margin/margin_ref)  ∈ [0,1]
    "w_conf_cost": 0.6, "w_conf_margin": 0.4,
}

NUMBERS = (1, 2, 3, 4, 5, 6, 7, 8, 9)   # cột j của ma trận cost ↔ số j+1


# ------------------------------------------------------------------ tầng 1

def _bgr_to_hsv(bgr: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """(N,3) BGR ∈ [0,1] → (H ∈ [0,360), S ∈ [0,1], V ∈ [0,1]).

    Công thức hexcone chuẩn; KHÔNG theo scale hue [0,180) của OpenCV.
    """
    b, g, r = bgr[..., 0], bgr[..., 1], bgr[..., 2]
    v = bgr.max(axis=-1)
    c = v - bgr.min(axis=-1)
    s = np.where(v > _EPS, c / np.maximum(v, _EPS), 0.0)
    safe_c = np.maximum(c, _EPS)
    h = np.select(
        [v == r, v == g],
        [(g - b) / safe_c % 6.0, (b - r) / safe_c + 2.0],
        default=(r - g) / safe_c + 4.0,
    ) * 60.0
    h = np.where(c > _EPS, h % 360.0, 0.0)
    return h, s, v


def _hue_circular_median(h: np.ndarray) -> float:
    """Median hue có xử lý vòng tròn — bi đỏ/nâu vắt qua mốc 0°/360°.

    Median thẳng trên {358°, 2°} ra ~180° (cyan!) — sai chết người với bi
    3/7. Cách làm: xoay cụm hue về quanh 180° theo hướng trung bình vector
    (atan2), lấy median trong hệ đã xoay, xoay ngược lại.
    """
    rad = np.deg2rad(h)
    mean_ang = np.rad2deg(np.arctan2(np.sin(rad).mean(), np.cos(rad).mean()))
    shift = 180.0 - mean_ang
    return float((np.median((h + shift) % 360.0) - shift) % 360.0)


def _crop(img: np.ndarray, box) -> np.ndarray | None:
    """Crop quanh tâm bbox, co ``crop_shrink`` kích thước; None nếu rỗng.

    bbox = (x1, y1, x2, y2) pixel trên CHÍNH ảnh đưa vào (không quy đổi
    scale ở đây — người gọi tự bảo đảm bbox đo trên cùng ảnh).
    """
    x1, y1, x2, y2 = (float(v) for v in box)
    cx, cy = (x1 + x2) / 2.0, (y1 + y2) / 2.0
    hw = (x2 - x1) * BALLID["crop_shrink"] / 2.0
    hh = (y2 - y1) * BALLID["crop_shrink"] / 2.0
    h_img, w_img = img.shape[:2]
    xa, xb = max(0, round(cx - hw)), min(w_img, round(cx + hw))
    ya, yb = max(0, round(cy - hh)), min(h_img, round(cy + hh))
    if xb - xa < 1 or yb - ya < 1:
        return None
    return img[ya:yb, xa:xb]


def _features(crop_bgr: np.ndarray, wb_scale: np.ndarray | None
              ) -> tuple[float, float, float, float]:
    """Crop BGR → (hue_median, sat_median, val_median, white_ratio).

    ``wb_scale`` (3,) nhân vào 3 kênh BGR (chuẩn hoá trắng theo cue) trước
    khi đọc màu; None = không WB. Median lấy trên pixel MÀU (không trắng);
    hue dùng median vòng tròn (xem `_hue_circular_median`).
    """
    px = crop_bgr.reshape(-1, 3).astype(np.float64)
    if wb_scale is not None:
        px = np.clip(px * wb_scale, 0.0, 255.0)
    h, s, v = _bgr_to_hsv(px / 255.0)
    white = (v >= BALLID["white_v_min"]) & (s <= BALLID["white_s_max"])
    white_ratio = float(white.mean())
    colored = ~white
    if colored.mean() < BALLID["min_color_frac"]:
        colored = np.ones_like(white)      # gần toàn trắng — median mọi pixel
    return (_hue_circular_median(h[colored]), float(np.median(s[colored])),
            float(np.median(v[colored])), white_ratio)


def _white_penalty(white_ratio: float, lo: float, hi: float) -> float:
    """Khoảng cách tuyến tính của white_ratio tới đoạn [lo, hi] (0 nếu trong)."""
    return max(0.0, lo - white_ratio, white_ratio - hi)


def _cost_vector(h: float, s: float, v: float, white_ratio: float) -> np.ndarray:
    """Đặc trưng một crop → vector cost 9 số (index j ↔ số j+1).

    Đây là CHỮ KÝ swap classifier (design §2): sau này thay bằng model trả
    −log p, tầng 2 giữ nguyên.
    """
    P = BALLID
    cost = np.empty(9, dtype=np.float64)
    for j, num in enumerate(NUMBERS):
        if num == 8:
            pen = _white_penalty(white_ratio, *P["white_range"]["eight"])
            cost[j] = (P["w8_val"] * max(0.0, v - P["v8_max"])
                       + P["w8_sat"] * max(0.0, s - P["s8_max"])
                       + P["w_white"] * pen)
            continue
        h_ref, s_ref, v_ref = P["palette"][num]
        d_h = abs(h - h_ref) % 360.0
        d_h = min(d_h, 360.0 - d_h) / 180.0
        rng = P["white_range"]["stripe" if num == 9 else "solid"]
        cost[j] = (P["w_hue"] * d_h
                   + P["w_sat"] * abs(s - s_ref)
                   + P["w_val"] * abs(v - v_ref)
                   + P["w_white"] * _white_penalty(white_ratio, *rng))
    return cost


def score_balls(img_bgr: np.ndarray, ball_boxes, cue_box=None
                ) -> tuple[np.ndarray, bool]:
    """Tầng 1: ảnh BGR + bbox các bi (+ bbox cue) → ma trận cost (N, 9).

    Trả (cost, wb): ``wb=True`` khi có chuẩn hoá trắng theo crop cue (cue
    hợp lệ trong ảnh); ``False`` khi không cue → đọc màu thô (design §2).
    Bbox hỏng/ngoài ảnh → hàng cost = +inf → tầng 2 tự trả ``None``.
    """
    img = np.asarray(img_bgr)
    if img.ndim != 3 or img.shape[2] != 3:
        raise ValueError(f"img_bgr phải là (H,W,3) BGR, nhận shape {img.shape}.")

    wb_scale = None
    if cue_box is not None:
        cue_crop = _crop(img, cue_box)
        if cue_crop is not None:
            med = np.median(cue_crop.reshape(-1, 3).astype(np.float64), axis=0)
            if np.all(med > _EPS):
                # cue trắng ám đèn thế nào thì scale kênh về trung tính thế ấy
                wb_scale = med.mean() / med
    wb = wb_scale is not None

    cost = np.full((len(ball_boxes), 9), np.inf, dtype=np.float64)
    for i, box in enumerate(ball_boxes):
        crop = _crop(img, box)
        if crop is None:
            continue
        cost[i] = _cost_vector(*_features(crop, wb_scale))
    return cost, wb


# ------------------------------------------------------------------ tầng 2

def assign_numbers(cost: np.ndarray) -> list[tuple[int | None, float]]:
    """Tầng 2: ma trận cost (N,9) → [(số 1–9 | None, number_conf)] mỗi bi.

    Greedy tất định (design §2, chốt thay Hungarian có chủ đích): duyệt cặp
    (bi, số) theo cost tăng dần, gán khi cả bi lẫn số còn trống; hoà cost →
    (index bi, index số) nhỏ trước — sort tuple bảo đảm tất định. Cost gán
    > ``cost_max`` (hoặc không còn số trống, tức N > 9) → ``None``.

    ``number_conf`` ∈ [0,1] nghịch biến cost + margin với ứng viên nhì
    TRONG HÀNG của bi đó (margin âm — bị ràng buộc ép nhận số không phải
    rẻ nhất — kẹp về 0 = "hai số tranh nhau, không chắc").
    """
    c = np.asarray(cost, dtype=np.float64)
    if c.ndim != 2 or c.shape[1] != 9:
        raise ValueError(f"cost phải là (N,9), nhận {c.shape}.")
    P = BALLID
    pairs = sorted(
        (c[i, j], i, j)
        for i in range(c.shape[0]) for j in range(9)
        if np.isfinite(c[i, j]) and c[i, j] <= P["cost_max"]
    )
    ball_of: dict[int, int] = {}     # j → i
    num_of: dict[int, int] = {}      # i → j
    for cost_ij, i, j in pairs:
        if i in num_of or j in ball_of:
            continue
        ball_of[j] = i
        num_of[i] = j

    out: list[tuple[int | None, float]] = []
    for i in range(c.shape[0]):
        j = num_of.get(i)
        if j is None:
            out.append((None, 0.0))
            continue
        others = np.delete(c[i], j)
        runner_up = float(others.min()) if np.isfinite(others).any() else np.inf
        margin = max(0.0, runner_up - float(c[i, j]))
        conf = (P["w_conf_cost"] * max(0.0, 1.0 - float(c[i, j]) / P["cost_max"])
                + P["w_conf_margin"] * min(1.0, margin / P["margin_ref"]))
        out.append((NUMBERS[j], round(min(1.0, conf), 3)))
    return out


def identify_balls(img_bgr: np.ndarray, ball_boxes, cue_box=None
                   ) -> tuple[list[tuple[int | None, float]], bool]:
    """Trọn gói hai tầng: ảnh + bbox → [(number|None, conf)] theo thứ tự
    ``ball_boxes``, kèm cờ ``wb``. Không bi nào → danh sách rỗng."""
    if not len(ball_boxes):
        return [], cue_box is not None
    cost, wb = score_balls(img_bgr, ball_boxes, cue_box)
    return assign_numbers(cost), wb