src/poolcoach_cv/segmentation.py

"""Shot segmentation — cắt video dài thành từng cú (lát A1, BRIEF 13/08/2026).

Tầng "định nghĩa cú" của design Analyzer đa cú (§4 PoolCoach_Analyzer_
DemoApp_Design.md): chạy trên TRACK THÔ toàn video (detect + nối NN, trước
verifier), cắt ranh giới cú theo đúng chu trình **mọi bi đứng yên → cue
chuyển động → mọi bi đứng yên**.

Hai tầng, cùng nếp ``broadcast.py``:

- ``segment_shots``  — lõi cắt ranh giới, THUẦN NUMPY (unit test chạy trên
  venv app, không cv2/torch). Input là track cue (``track_cue`` của
  broadcast) + detection KHÔNG-cue per frame.
- ``segment_video``  — front-end trọn gói cho MỘT video dài: decode PTS thật
  + img_diff, YOLO detect từng frame, track cue, rồi ``segment_shots``.
  Import cv2 LƯỜI trong thân hàm (nếp analyze_clip).

Ngưỡng VẬT LÝ dùng NGUYÊN bộ đã chốt 11/08 (BRIEF A1: "dùng nguyên, ĐỪNG
chế ngưỡng mới") — tất cả import từ ``broadcast.py``, không định nghĩa lại:
đứng yên ``V_STILL_MPS`` 55 mm/s (p99 bi tĩnh P0); onset ``V_INIT_MPS``
0.25 m/s × ``V_INIT_RUN`` 3 bước (measure_p0); gating track ``VMAX_MPS``
11 m/s theo dt timestamp thật. Các hằng MỚI ở đây đều là hằng CẤU TRÚC
(cửa sổ thời gian, đệm biên) — không phải ngưỡng đo vật lý mới; nguồn ghi
tại chỗ từng hằng, khai trong HANDOFF A1 để Cowork duyệt.

Quy ước output: mỗi cú một dict, thời gian TUYỆT ĐỐI theo PTS video —
``t_start_s``/``t_end_s`` là biên ĐÃ đệm (cắt clip theo cặp này);
``t_onset_s``/``t_settle_s`` là mốc thô (cue bắt đầu chạy / mọi bi yên trở
lại) để đối chiếu đếm tay. ``status`` ∈ {"ok", "error"}; cú "error" vẫn
NẰM TRONG danh sách kèm ``reason`` tiếng Việt (scope §5 design: đổi cảnh
giữa cú phải báo tử tế, không nuốt im).
"""

from __future__ import annotations

import math
from pathlib import Path

import numpy as np

from poolcoach_cv.broadcast import (ANALYZE_MAX_DUR_S, BALL_CLASSES,
                                    DETECT_BATCH, GATE_SLACK_M, IN_TOL_M,
                                    TABLE_L_M, TABLE_W_M, V_INIT_MPS,
                                    V_INIT_RUN, V_STILL_MPS, VMAX_MPS,
                                    find_dup_frames, track_cue)
from poolcoach_cv.homography import TableHomography, orient_corners

# ---------------------------------------------------- hằng CẤU TRÚC (mới)
# Không hằng nào dưới đây là ngưỡng ĐO vật lý — ngưỡng đo dùng nguyên bộ
# chốt 11/08 import từ broadcast (xem docstring module).

STILL_WIN_S = 0.5     # cửa sổ CHORD đo "đứng yên": bi tĩnh jitter ~2mm/frame
                      # (p99 apparent speed 55mm/s ⇒ ~1.8mm ở dt 33ms) nên
                      # chord 2 đầu cửa sổ 0.5s của bi tĩnh ≪ 55mm/s ×
                      # 0.5s = 27.5mm — chord giết nhiễu per-step mà không
                      # cần ngưỡng mới; bi lăn chậm 0.1 m/s đi 50mm > 27.5mm
                      # vẫn bị bắt là đang chạy.
BLIND_MAX_S = 1.0     # không thấy BI NÀO (cue lẫn thường) liên tục quá mức
                      # này = mất cảnh (replay/đổi góc máy) — cú đang mở
                      # đóng lại thành "error", không cố nối track xuyên
                      # cảnh cắt (scope §5 design).
PAD_PRE_S = 0.5       # đệm TRƯỚC onset khi cắt clip: analyze cần cửa sổ
                      # tĩnh trước cú — ``_static_balls`` đòi ≥ 3 detection
                      # mỗi cụm (STATIC_MIN_DETS) và ``balls_init`` lấy
                      # frame đầu; 0.5s ≈ 15 frame @30fps là đủ dư. Clamp
                      # về t_settle cú trước nên không bao giờ nuốt đuôi
                      # chuyển động của cú trước.
PAD_POST_S = 0.5      # đệm SAU t_settle: giữ trọn pha lăn chậm cuối cú cho
                      # analyze; clamp về onset cú sau.
SHOT_MAX_S = ANALYZE_MAX_DUR_S
                      # cú mở quá mức trần clip analyze (30s) mà chưa đóng
                      # được = ranh giới hỏng (track đứt kéo dài / bi không
                      # bao giờ yên) → đóng thành "error" thay vì đẻ clip
                      # mà /api/analyze từ chối.

LY_DO_MAT_CANH = ("mất cảnh giữa cú — không thấy bi nào trên bàn quá "
                  f"{BLIND_MAX_S:g}s (video đổi góc/replay); không nối "
                  "track xuyên cảnh cắt")
LY_DO_QUA_DAI = (f"cú mở quá {SHOT_MAX_S:g}s mà mọi bi chưa đứng yên trở "
                 "lại — ranh giới không đóng được (track đứt kéo dài hoặc "
                 "cảnh không phải một cú trọn)")
LY_DO_CUT_CUOI_VIDEO = ("video kết thúc giữa cú — không thấy lúc mọi bi "
                        "đứng yên trở lại")


# ------------------------------------------------- tracklet bi (không-cue)

def _build_tracklets(frames: list[tuple[str, float]],
                     others_by_frame: dict[str, list[tuple[float, float]]]
                     ) -> list[list[tuple[float, float, float]]]:
    """Nối detection KHÔNG-cue thành tracklet bằng NN per-step + gating
    VMAX theo dt thật (đúng luật track_cue — không thêm luật mới). Greedy
    theo khoảng cách tăng dần; det không nối được mở tracklet mới. Trả
    list tracklet, mỗi cái là list (t, x, y).

    Mục đích DUY NHẤT là đo "mọi bi đứng yên" — không phải BallID, không
    cần bền qua che khuất: tracklet đứt thì mở cái mới, chord đo trên từng
    mảnh vẫn nói đúng chuyện bi đang chạy hay đứng.
    """
    tracklets: list[list[tuple[float, float, float]]] = []
    open_idx: list[int] = []          # tracklet còn "sống" ở frame trước
    prev_t: float | None = None
    for name, t in frames:
        dets = others_by_frame.get(name, [])
        if not dets:
            # frame không có det không-cue: tracklet nào cũng khép — bi
            # mất dấu một frame sẽ mở tracklet mới, đủ cho mục đích đo.
            open_idx = []
            prev_t = t
            continue
        if not open_idx or prev_t is None:
            open_idx = []
            for x, y in dets:
                tracklets.append([(t, x, y)])
                open_idx.append(len(tracklets) - 1)
            prev_t = t
            continue
        dt = max(t - prev_t, 1e-3)
        gate = VMAX_MPS * dt + GATE_SLACK_M
        pairs = []                    # (dist, det_i, tracklet_j)
        for i, (x, y) in enumerate(dets):
            for j in open_idx:
                lx, ly = tracklets[j][-1][1], tracklets[j][-1][2]
                d = math.hypot(x - lx, y - ly)
                if d <= gate:
                    pairs.append((d, i, j))
        pairs.sort(key=lambda p: p[0])
        used_det: set[int] = set()
        used_trk: set[int] = set()
        next_open: list[int] = []
        for d, i, j in pairs:
            if i in used_det or j in used_trk:
                continue
            used_det.add(i)
            used_trk.add(j)
            tracklets[j].append((t, dets[i][0], dets[i][1]))
            next_open.append(j)
        for i, (x, y) in enumerate(dets):
            if i not in used_det:
                tracklets.append([(t, x, y)])
                next_open.append(len(tracklets) - 1)
        open_idx = next_open
        prev_t = t
    return tracklets


def _moving_times(tracklets, win_s: float) -> np.ndarray:
    """Mốc thời gian mà MỘT bi không-cue nào đó đang CHẠY: điểm tracklet có
    chord-speed nhìn lùi ``win_s`` vượt V_STILL. Trả mảng t đã sort."""
    out: list[float] = []
    for trk in tracklets:
        if len(trk) < 2:
            continue
        ts = np.array([p[0] for p in trk])
        xy = np.array([[p[1], p[2]] for p in trk])
        for i in range(1, len(trk)):
            j = int(np.searchsorted(ts, ts[i] - win_s, side="left"))
            j = min(j, i - 1)
            span = float(ts[i] - ts[j])
            if span <= 1e-6:
                continue
            chord = float(np.hypot(*(xy[i] - xy[j])))
            if chord / span > V_STILL_MPS:
                out.append(float(ts[i]))
    return np.asarray(sorted(out))


# ------------------------------------------------------------- lõi cắt cú

def segment_shots(rows: list[dict],
                  others: list[dict] | None = None) -> dict:
    """Cắt ranh giới cú trên track thô toàn video.

    ``rows``: track cue per frame (``track_cue`` broadcast — frame_file,
    t_s, covered, table_x_m, table_y_m, img_diff; str từ CSV hay số đều
    nhận). ``others``: detection KHÔNG-cue per frame ``{frame_file, t_s,
    x_m, y_m}`` — cần cho vế "MỌI bi đứng yên"; None/[] thì chỉ còn cue
    làm chứng (vẫn chạy, kèm warning).

    Trả dict JSON-thuần::

        {"shots": [{"idx", "t_start_s", "t_end_s", "t_onset_s",
                    "t_settle_s", "status", "reason"}...],
         "warnings": [...], "n_frames", "n_dup_frames", "duration_s"}

    Chu trình trạng thái: chờ MỌI bi đứng yên (ARMED) → cue chạy
    ``V_INIT_RUN`` bước liên tiếp > ``V_INIT_MPS`` (onset, đúng luật
    ``_find_motion_start``) → mở cú → mọi bi yên lại trọn cửa sổ
    ``STILL_WIN_S`` (settle) → đóng cú, quay lại ARMED. Mất cảnh
    (``BLIND_MAX_S`` không thấy bi nào) hay cú mở quá ``SHOT_MAX_S`` →
    đóng thành "error" kèm lý do, KHÔNG nuốt im.
    """
    if not rows:
        raise ValueError("Track rỗng — không có frame nào để cắt cú.")

    names = [str(r["frame_file"]) for r in rows]
    t_all = np.array([float(r["t_s"]) for r in rows])
    covered = np.array([str(r.get("covered", "0")) in ("1", "True", "true")
                        for r in rows])
    diffs = [float(r.get("img_diff", -1.0)) for r in rows]
    dups = find_dup_frames(names, diffs)
    keep = [i for i in range(len(rows)) if names[i] not in dups]

    frames = [(names[i], float(t_all[i])) for i in keep]
    t_frames = np.array([t for _n, t in frames])

    # --- cue: chuỗi điểm covered (đã bỏ frame trùng) + speed per-step
    cue_idx = [i for i in keep if covered[i]]
    t_cue = t_all[cue_idx]
    xy_cue = (np.array([[float(rows[i]["table_x_m"]),
                         float(rows[i]["table_y_m"])] for i in cue_idx])
              if cue_idx else np.empty((0, 2)))
    if len(t_cue) >= 2:
        dt_c = np.diff(t_cue)
        dxy_c = np.diff(xy_cue, axis=0)
        v_cue = np.hypot(dxy_c[:, 0], dxy_c[:, 1]) / np.maximum(dt_c, 1e-6)
        t_cue_step = t_cue[1:]            # mốc CUỐI mỗi bước
    else:
        v_cue = np.empty(0)
        t_cue_step = np.empty(0)

    # chord-speed nhìn lùi của cue (cùng thước với bi thường) cho vế settle
    cue_moving_t: list[float] = []
    for i in range(1, len(t_cue)):
        j = int(np.searchsorted(t_cue, t_cue[i] - STILL_WIN_S, side="left"))
        j = min(j, i - 1)
        span = float(t_cue[i] - t_cue[j])
        if span <= 1e-6:
            continue
        if float(np.hypot(*(xy_cue[i] - xy_cue[j]))) / span > V_STILL_MPS:
            cue_moving_t.append(float(t_cue[i]))
    cue_moving = np.asarray(cue_moving_t)

    # --- bi không-cue: tracklet + mốc "đang chạy"
    others = others or []
    others_by_frame: dict[str, list[tuple[float, float]]] = {}
    for o in others:
        fname = str(o["frame_file"])
        if fname in dups:
            continue
        others_by_frame.setdefault(fname, []).append(
            (float(o["x_m"]), float(o["y_m"])))
    tracklets = _build_tracklets(frames, others_by_frame)
    others_moving = _moving_times(tracklets, STILL_WIN_S)

    # --- frame mù: không thấy bi NÀO (cue lẫn thường)
    blind = np.array([not covered[i] and not others_by_frame.get(names[i])
                      for i in keep])

    # mốc "có dữ liệu": mọi thời điểm một chord đánh giá được (điểm thứ ≥2
    # của cue hoặc của một tracklet) — cửa sổ không có mốc nào thì "yên" là
    # KHÔNG BIẾT, không phải yên. Precompute một mảng sort để still_at chạy
    # O(log n) mỗi frame (quét video dài hàng nghìn frame).
    data_list = [float(v) for v in t_cue[1:]]
    for trk in tracklets:
        data_list.extend(p[0] for p in trk[1:])
    data_times = np.asarray(sorted(data_list))

    def any_moving(t_a: float, t_b: float) -> bool:
        """Có bi nào (cue hay thường) mang cờ 'đang chạy' trong (t_a, t_b]?"""
        for arr in (cue_moving, others_moving):
            lo = int(np.searchsorted(arr, t_a, side="right"))
            hi = int(np.searchsorted(arr, t_b, side="right"))
            if hi > lo:
                return True
        return False

    def cue_seen_moving_before(t: float) -> bool:
        """Cue mất dấu: bước cuối trước ``t`` còn > V_INIT? (đang bay thì
        mất det — chưa được kết luận 'mọi bi yên' chỉ vì vắng mặt nó)."""
        k = int(np.searchsorted(t_cue_step, t, side="right")) - 1
        return k >= 0 and float(v_cue[k]) > V_INIT_MPS

    def still_at(t_end_win: float, in_shot: bool) -> bool:
        t_a = t_end_win - STILL_WIN_S
        if t_a < float(t_frames[0]) - 1e-9:
            return False
        lo = int(np.searchsorted(data_times, t_a, side="left"))
        hi = int(np.searchsorted(data_times, t_end_win, side="right"))
        if hi <= lo:
            return False
        if any_moving(t_a, t_end_win):
            return False
        if in_shot:
            # cue vắng mặt trọn cửa sổ mà lần cuối thấy nó còn đang bay →
            # chưa yên (bài học shot_07: người che cue giữa cú)
            c_lo = int(np.searchsorted(t_cue, t_a, side="left"))
            c_hi = int(np.searchsorted(t_cue, t_end_win, side="right"))
            if c_hi - c_lo < 2 and cue_seen_moving_before(t_a):
                return False
        return True

    # --- onset ứng viên: V_INIT_RUN bước cue liên tiếp > V_INIT_MPS
    # (đúng luật _find_motion_start của measure_p0/broadcast; mốc onset =
    # điểm ĐẦU của run). Bước vắt qua đứt track dài hơn cửa sổ không được
    # tính "liên tiếp" — re-capture sau coast dài cho speed chord to giả.
    onsets: list[float] = []
    run = 0
    for k in range(len(v_cue)):
        if (t_cue[k + 1] - t_cue[k]) > STILL_WIN_S:
            run = 0
        run = run + 1 if v_cue[k] > V_INIT_MPS else 0
        if run == V_INIT_RUN:
            onsets.append(float(t_cue[k + 1 - V_INIT_RUN]))
    onsets_arr = np.asarray(onsets)

    # --- máy trạng thái quét theo frame đã bỏ trùng
    shots: list[dict] = []
    warnings: list[str] = []
    state = "idle"                    # idle → armed → open
    armed_at: float | None = None     # mốc "mọi bi đã yên" gần nhất
    t_onset: float | None = None
    blind_from: float | None = None

    def close(t_settle: float, status: str, reason: str | None) -> None:
        nonlocal state, armed_at, t_onset
        shots.append({"t_onset_s": float(t_onset),
                      "t_settle_s": float(t_settle),
                      "status": status, "reason": reason})
        if status == "ok":
            # bàn ĐANG yên (chính là điều kiện vừa đóng cú) → armed luôn,
            # không quay về idle: cú kế bắn nhanh ngay sau settle vẫn bắt.
            state, armed_at = "armed", t_settle
        else:
            state, armed_at = "idle", None
        t_onset = None

    for fi, (name, t) in enumerate(frames):
        if blind[fi]:
            if blind_from is None:
                blind_from = t
            if t - blind_from > BLIND_MAX_S:
                if state == "open":
                    close(blind_from, "error", LY_DO_MAT_CANH)
                # mất cảnh xoá cả bằng chứng "đã yên" — phải yên lại từ đầu
                state, armed_at = "idle", None
            continue
        blind_from = None

        if state == "idle" and still_at(t, in_shot=False):
            state = "armed"
            armed_at = t - STILL_WIN_S
        if state == "armed":
            lo = int(np.searchsorted(onsets_arr, armed_at - 1e-9,
                                     side="left"))
            hi = int(np.searchsorted(onsets_arr, t + 1e-9, side="right"))
            if hi > lo:
                t_onset = float(onsets_arr[lo])
                state = "open"
                continue
        if state == "open":
            if t - t_onset > SHOT_MAX_S:
                close(t, "error", LY_DO_QUA_DAI)
            elif t - t_onset > STILL_WIN_S and still_at(t, in_shot=True):
                close(t - STILL_WIN_S, "ok", None)

    if state == "open":
        close(float(t_frames[-1]), "error", LY_DO_CUT_CUOI_VIDEO)

    # --- đệm biên + clamp vào video và cú lân cận
    t0, t1 = float(t_frames[0]), float(t_frames[-1])
    for i, s in enumerate(shots):
        prev_end = shots[i - 1]["t_settle_s"] if i > 0 else t0
        next_on = shots[i + 1]["t_onset_s"] if i + 1 < len(shots) else t1
        s["idx"] = i + 1
        s["t_start_s"] = round(max(s["t_onset_s"] - PAD_PRE_S, prev_end, t0), 3)
        s["t_end_s"] = round(min(s["t_settle_s"] + PAD_POST_S, next_on, t1), 3)
        s["t_onset_s"] = round(s["t_onset_s"], 3)
        s["t_settle_s"] = round(s["t_settle_s"], 3)

    if not others:
        warnings.append("Không có detection bi thường nào — vế 'mọi bi đứng "
                        "yên' chỉ còn cue ball làm chứng.")
    return {"shots": shots,
            "warnings": warnings,
            "n_frames": len(rows),
            "n_dup_frames": len(dups),
            "duration_s": round(float(t_all[-1] - t_all[0]), 3)
            if len(t_all) > 1 else 0.0}


# ------------------------------------------ front-end video (cv2/YOLO lười)

def segment_video(video_path: str | Path, corners_px, model, *,
                  op_conf: float, imgsz: int = 640, device=None,
                  max_dur_s: float | None = None, progress=None) -> dict:
    """MỘT video dài → detect + track toàn video → danh sách cú.

    Cùng giao ước ``analyze_clip`` (corners 4 pocket pixel, model YOLO đã
    load, ``progress(stage, frac)``, ValueError cho lỗi input) nhưng KHÔNG
    trần 30s — trần video do route quyết (``max_dur_s``). Decode trong bộ
    nhớ, không ghi frame ra đĩa (ràng buộc bản quyền, nếp BRIEF 24 #7).

    Trả dict của ``segment_shots`` cộng ``fps_nominal`` và
    ``t_first_s``/``t_last_s`` (PTS tuyệt đối — clip per cú cắt theo đúng
    thước này).
    """
    import cv2

    corners = np.asarray(corners_px, dtype=np.float64)
    if corners.shape != (4, 2):
        raise ValueError("corners phải là 4 điểm [x, y] pixel — chấm đủ "
                         "4 pocket góc.")
    th = TableHomography(orient_corners(corners),
                         table_w=TABLE_W_M, table_l=TABLE_L_M)

    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        raise ValueError("Không mở được video — cần MP4/MOV (H.264) đọc được.")
    try:
        fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
        n_est = cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0
        if (max_dur_s and fps > 0 and n_est > 0
                and n_est / fps > max_dur_s + 0.5):
            raise ValueError(f"Video dài {n_est / fps / 60:.1f} phút — giới "
                             f"hạn {max_dur_s / 60:g} phút cho một lần quét.")

        ci = th.corners_px.astype(int)
        x0, x1 = ci[:, 0].min(), ci[:, 0].max()
        y0, y1 = ci[:, 1].min(), ci[:, 1].max()

        meta: list[tuple[str, float]] = []
        diffs: dict[str, float] = {}
        dets_by_frame: dict[str, list[dict]] = {}
        other_dets: list[dict] = []
        prev_crop = None
        t_first = None
        batch_imgs: list = []
        batch_names: list[str] = []
        total = int(n_est) if n_est > 0 else 0

        def flush_batch() -> None:
            if not batch_imgs:
                return
            kw = {"device": device} if device is not None else {}
            results = model.predict(batch_imgs, conf=op_conf, imgsz=imgsz,
                                    verbose=False, **kw)
            for name, res in zip(batch_names, results):
                names_map = res.names
                for b in res.boxes:
                    cls = names_map[int(b.cls)]
                    if cls not in BALL_CLASSES:
                        continue
                    bx1, by1, bx2, by2 = (float(v) for v in b.xyxy[0])
                    cx, cy = (bx1 + bx2) / 2.0, (by1 + by2) / 2.0
                    tx, ty = th.px_to_table((cx, cy))
                    if not (-IN_TOL_M <= tx <= TABLE_W_M + IN_TOL_M
                            and -IN_TOL_M <= ty <= TABLE_L_M + IN_TOL_M):
                        continue
                    if cls == "Cue":
                        dets_by_frame.setdefault(name, []).append(
                            {"conf": float(b.conf), "table_x_m": float(tx),
                             "table_y_m": float(ty)})
                    else:
                        other_dets.append({"frame_file": name,
                                           "t_s": meta[int(name)][1],
                                           "x_m": float(tx),
                                           "y_m": float(ty)})
            batch_imgs.clear()
            batch_names.clear()

        while True:
            ok = cap.grab()
            if not ok:
                break
            t = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.0
            if t_first is None:
                t_first = t
            if max_dur_s and t - t_first > max_dur_s + 0.5:
                raise ValueError(f"Video dài quá {max_dur_s / 60:g} phút "
                                 f"(container khai sai số frame).")
            ok, frame = cap.retrieve()
            if not ok:
                break
            name = f"{len(meta):05d}"
            crop = frame[max(0, y0):max(1, y1), max(0, x0):max(1, x1)]
            gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY).astype(np.int16)
            diffs[name] = (float(np.abs(gray - prev_crop).mean())
                           if prev_crop is not None else -1.0)
            prev_crop = gray
            meta.append((name, t))
            batch_imgs.append(frame)
            batch_names.append(name)
            if len(batch_imgs) >= DETECT_BATCH:
                flush_batch()
                if progress is not None:
                    frac = min(len(meta) / total, 1.0) if total else 0.0
                    progress("detect", 0.02 + 0.9 * frac)
        flush_batch()
    finally:
        cap.release()

    if len(meta) < 5:
        raise ValueError("Video quá ngắn hoặc không đọc được frame nào.")

    if progress is not None:
        progress("segment", 0.95)
    rows = track_cue(meta, dets_by_frame, diffs)
    out = segment_shots(rows, others=other_dets)
    out["fps_nominal"] = round(fps, 3)
    out["t_first_s"] = round(float(meta[0][1]), 3)
    out["t_last_s"] = round(float(meta[-1][1]), 3)
    return out