scripts/broadcast/measure_p0.py

"""Bước 4 P0 (BRIEF 11/08/2026): đo số gate G0 từ track đã sinh ở bước 3.

Đọc per-shot ``track.csv`` + ``detections.csv`` + ``meta.json``, xuất CSV ra
ROOT ``D:\\Khoa luan`` theo nếp eval, tiền tố ``bb9_p0_``:

- ``bb9_p0_shots.csv``    — per shot: tỷ lệ frame detect cue, coverage theo
  THỜI LƯỢNG (số gate G0), max gap, thời điểm cue bắt đầu chạy, số gãy khúc.
- ``bb9_p0_noise.csv``    — noise vị trí (mm) per-ball: các bi đứng yên trong
  cửa sổ trước khi cue chạy → std/phân vị sau homography.
- ``bb9_p0_segments.csv`` — verifier vật lý bản thô: giữa hai gãy khúc fit
  đoạn thẳng (PCA) + giảm tốc tuyến tính; residual mm/phân vị. CHỈ fit-và-báo.
- ``bb9_p0_summary.csv``  — bảng tổng + verdict G0.
- ``bb9_p0_plots\\``       — speed trace + path per shot, histogram tổng.

Ngưỡng "đứng yên" và "gãy khúc" là OUTPUT (đề xuất từ percentile đo thật,
kèm biểu đồ — Cowork chốt vào design doc §5), KHÔNG phải input chỉnh tay.
Cửa sổ đứng yên cắt bằng ngưỡng khởi động bảo thủ V_INIT (hằng số, ghi rõ)
— không phụ thuộc số đề xuất để tránh vòng lặp gà-trứng.

Ghi chú tay per-shot (cue có chạm băng/bi sau chạm đầu → spin identifiable?)
đọc từ ``<root>/notes.json`` nếu có: {"1": {"contact_after_first": true,
"comment": "..."}, ...} — thống kê tỷ lệ spin-unidentifiable (rủi ro §10).

    D:\\Khoa luan\\poolcoach-cv-env\\Scripts\\python.exe \
        scripts/broadcast/measure_p0.py
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]  # poolcoach-rl/
sys.path.insert(0, str(ROOT / "src"))

import numpy as np  # noqa: E402

PILOT_ROOT = ROOT / "datasets" / "bb9_pilot"
OUT_ROOT = ROOT.parent                      # D:\Khoa luan — nếp eval CSV ra root
BALL_CLASSES = {"Black", "Cue", "Solid", "Striped"}

V_INIT = 0.25        # m/s — ngưỡng BẢO THỦ chỉ để cắt cửa sổ đứng yên
V_INIT_RUN = 3       # số bước liên tiếp vượt V_INIT mới tính là "bắt đầu chạy"
STILL_MARGIN = 2     # bỏ 2 frame sát lúc chạy khỏi cửa sổ đứng yên
CLUSTER_R_M = 0.06   # bán kính gom det đứng yên về cùng một bi (~2R bi)
MIN_STILL_FRAMES = 5
GAP_SPLIT_S = 0.15   # track đứt quá mức này → cắt đoạn fit
MIN_SEG_PTS = 6
PCTS = (50, 90, 95, 99)


def read_csv(path: Path) -> list[dict]:
    with open(path, encoding="utf-8") as f:
        return list(csv.DictReader(f))


def write_csv(path: Path, rows: list[dict]) -> None:
    if not rows:
        return
    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)


def pct_dict(a: np.ndarray, prefix: str, scale: float = 1.0,
             nd: int = 2) -> dict:
    out = {}
    for p in PCTS:
        out[f"{prefix}_p{p}"] = round(float(np.percentile(a, p)) * scale, nd)
    out[f"{prefix}_max"] = round(float(a.max()) * scale, nd)
    return out


def find_dup_frames(track: list[dict]) -> set[str]:
    """Frame nhân đôi (stream 30fps từ nguồn 25fps → ~1/6 frame gần trùng).

    Luật CỤC BỘ trên cột ``img_diff`` (track_p0): frame là bản trùng nếu
    diff của nó < 0.3 × median diff của ≤6 frame lân cận VÀ lân cận đang
    thực sự chuyển động (median > 0.08 — tránh cờ giả trong cảnh tĩnh,
    nơi mọi diff đều ~nhiễu encoder).
    """
    names = [r["frame_file"] for r in track]
    diffs = [float(r.get("img_diff", -1)) for r in track]
    dups: set[str] = set()
    for i in range(1, len(diffs)):
        if diffs[i] < 0:
            continue
        win = [diffs[j] for j in range(max(1, i - 3), min(len(diffs), i + 4))
               if j != i and diffs[j] >= 0]
        if not win:
            continue
        med = float(np.median(win))
        if med > 0.08 and diffs[i] < 0.3 * med:
            dups.add(names[i])
    return dups


def cue_speed_series(track: list[dict], skip: frozenset = frozenset()):
    """(t, x, y) các frame covered → (t_mid, speed m/s, heading rad).

    ``skip``: frame nhân đôi cần loại khỏi chuỗi vận tốc (dt tự cộng dồn
    sang frame thật kế tiếp — timestamp thật vẫn đúng).
    """
    pts = [(float(r["t_s"]), float(r["table_x_m"]), float(r["table_y_m"]))
           for r in track
           if r["covered"] == "1" and r["frame_file"] not in skip]
    t = np.array([p[0] for p in pts])
    xy = np.array([[p[1], p[2]] for p in pts])
    if len(pts) < 2:
        return t, xy, np.array([]), np.array([]), np.array([])
    dt = np.diff(t)
    dxy = np.diff(xy, axis=0)
    dist = np.hypot(dxy[:, 0], dxy[:, 1])
    speed = dist / np.maximum(dt, 1e-6)
    heading = np.arctan2(dxy[:, 1], dxy[:, 0])
    t_mid = t[:-1] + dt / 2
    return t, xy, t_mid, speed, heading


def find_motion_start(t_mid: np.ndarray, speed: np.ndarray) -> float | None:
    run = 0
    for i, v in enumerate(speed):
        run = run + 1 if v > V_INIT else 0
        if run >= V_INIT_RUN:
            return float(t_mid[i - V_INIT_RUN + 1])
    return None


def noise_for_shot(dets: list[dict], t0: float, t_still_end: float,
                   shot: str) -> tuple[list[dict], np.ndarray, np.ndarray]:
    """Cluster det đứng yên per-ball → noise mm + apparent speed pooled."""
    still = [d for d in dets
             if d["at_op"] == "1" and d["in_table"] == "1"
             and d["cls"] in BALL_CLASSES
             and t0 <= float(d["t_s"]) <= t_still_end]
    if not still:
        return [], np.array([]), np.array([])
    frames = sorted({d["frame_file"] for d in still})
    n_frames = len(frames)
    first = [d for d in still if d["frame_file"] == frames[0]]
    clusters = [{"cls": d["cls"], "pts": [], "ts": []} for d in first]
    anchors = np.array([[float(d["table_x_m"]), float(d["table_y_m"])]
                        for d in first])
    for d in still:
        p = np.array([float(d["table_x_m"]), float(d["table_y_m"])])
        dist = np.hypot(*(anchors - p).T)
        j = int(dist.argmin())
        if dist[j] <= CLUSTER_R_M:
            clusters[j]["pts"].append(p)
            clusters[j]["ts"].append(float(d["t_s"]))

    rows, devs_all, speeds_all = [], [], []
    for j, c in enumerate(clusters):
        if len(c["pts"]) < max(MIN_STILL_FRAMES, int(0.5 * n_frames)):
            continue
        pts = np.array(c["pts"])
        mean = pts.mean(axis=0)
        dev = np.hypot(*(pts - mean).T)          # m, so voi tam trung binh
        order = np.argsort(c["ts"])
        ts = np.array(c["ts"])[order]
        ps = pts[order]
        dt = np.diff(ts)
        ok = dt > 1e-6
        sp = np.hypot(*np.diff(ps, axis=0).T)[ok] / dt[ok]
        row = {"shot": shot, "ball_idx": j, "cls": c["cls"],
               "n_dets": len(pts), "window_frames": n_frames,
               "std_x_mm": round(float(pts[:, 0].std()) * 1000, 2),
               "std_y_mm": round(float(pts[:, 1].std()) * 1000, 2)}
        row.update(pct_dict(dev, "r_mm", scale=1000.0))
        rows.append(row)
        devs_all.append(dev)
        speeds_all.append(sp)
    devs = np.concatenate(devs_all) if devs_all else np.array([])
    spds = np.concatenate(speeds_all) if speeds_all else np.array([])
    return rows, devs, spds


def split_segments(t: np.ndarray, xy: np.ndarray, heading: np.ndarray,
                   speed: np.ndarray, t_mid: np.ndarray, v_still: float,
                   kink_deg: float):
    """Cắt track cue thành đoạn giữa các gãy khúc / khoảng đứt / dừng."""
    moving = speed > max(v_still, V_INIT / 2)
    dtheta = np.abs(np.degrees(np.diff(np.unwrap(heading))))
    kinks = []                      # index vao mang diem (xy)
    for i in range(len(dtheta)):
        if moving[i] and moving[i + 1] and dtheta[i] > kink_deg:
            kinks.append(i + 1)     # diem giua 2 buoc doi huong

    cut = set(kinks)
    for i in range(len(t) - 1):
        if t[i + 1] - t[i] > GAP_SPLIT_S:
            cut.add(i + 1)
    for i in range(len(speed)):
        if not moving[i]:
            cut.add(i + 1)

    segs, start = [], 0
    for i in sorted(cut) + [len(t)]:
        if i - start >= MIN_SEG_PTS:
            segs.append((start, i))
        start = max(start, i)
    return segs, kinks, dtheta, moving


def fit_segment(t: np.ndarray, xy: np.ndarray) -> dict:
    """PCA line fit (residual vuông góc, mm) + giảm tốc tuyến tính."""
    mean = xy.mean(axis=0)
    X = xy - mean
    _, _, Vt = np.linalg.svd(X, full_matrices=False)
    d = Vt[0]
    perp = np.abs(X @ np.array([-d[1], d[0]]))    # m
    dt = np.diff(t)
    sp = np.hypot(*np.diff(xy, axis=0).T) / np.maximum(dt, 1e-6)
    t_mid = t[:-1] + dt / 2
    if len(sp) >= 2:
        A = np.vstack([t_mid - t_mid[0], np.ones_like(t_mid)]).T
        coef, *_ = np.linalg.lstsq(A, sp, rcond=None)
        v_resid = sp - A @ coef
        decel = -float(coef[0])
    else:
        v_resid = np.array([0.0])
        decel = float("nan")
    out = {"n_pts": len(t), "dur_s": round(float(t[-1] - t[0]), 3),
           "decel_mps2": round(decel, 3),
           "v_mean_mps": round(float(sp.mean()), 3) if len(sp) else 0.0}
    out.update(pct_dict(perp, "line_resid_mm", scale=1000.0))
    out["v_resid_rms_mps"] = round(float(np.sqrt((v_resid ** 2).mean())), 4)
    return out


def main() -> None:
    ap = argparse.ArgumentParser(description="Do so gate G0 (P0)")
    ap.add_argument("--root", type=Path, default=PILOT_ROOT)
    ap.add_argument("--out-root", type=Path, default=OUT_ROOT)
    ap.add_argument("--only", type=int, nargs="*", default=None)
    args = ap.parse_args()

    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plots = args.out_root / "bb9_p0_plots"
    plots.mkdir(exist_ok=True)

    shots = sorted(d for d in args.root.glob("shot_*") if d.is_dir()
                   if (d / "track.csv").exists())
    if args.only:
        want = {f"shot_{i:02d}" for i in args.only}
        shots = [d for d in shots if d.name in want]
    if not shots:
        sys.exit("[ERROR] Khong co shot nao co track.csv — chay track_p0.py truoc.")

    notes_path = args.root / "notes.json"
    notes = (json.loads(notes_path.read_text(encoding="utf-8"))
             if notes_path.exists() else {})

    shot_rows, noise_rows, seg_rows = [], [], []
    all_dev, all_still_sp, all_dtheta_smooth = [], [], []

    # ---- pass 1: noise + still-speed de xuat v_still --------------------
    per_shot_cache = {}
    for d in shots:
        track = read_csv(d / "track.csv")
        dets = read_csv(d / "detections.csv") if (d / "detections.csv").exists() else []
        t_all = np.array([float(r["t_s"]) for r in track])
        dups = frozenset(find_dup_frames(track))
        t, xy, t_mid, speed, heading = cue_speed_series(track, skip=dups)
        motion_start = find_motion_start(t_mid, speed) if len(speed) else None
        dt_med = float(np.median(np.diff(t_all))) if len(t_all) > 1 else 1 / 30
        t_still_end = (motion_start - STILL_MARGIN * dt_med
                       if motion_start is not None else t_all[0] + 5 * dt_med)
        nrows, dev, sp = noise_for_shot(dets, t_all[0], t_still_end, d.name)
        noise_rows += nrows
        if len(dev):
            all_dev.append(dev)
        if len(sp):
            all_still_sp.append(sp)
        per_shot_cache[d.name] = (track, dets, t_all, t, xy, t_mid, speed,
                                  heading, motion_start, dt_med, t_still_end,
                                  dups)

    still_sp = np.concatenate(all_still_sp) if all_still_sp else np.array([0.0])
    v_still_prop = float(np.percentile(still_sp, 99))

    # ---- pass 2: kink + segments voi v_still de xuat --------------------
    # Nguong gay khuc de xuat: p95 cua |dtheta| khi dang chay TRON —
    # lay tu phan phoi thuc te, bao cao percentile de Cowork chot.
    dtheta_pool, dtheta_pool_raw = [], []
    for d in shots:
        (track, _dets, _t_all, t, xy, t_mid, speed, heading,
         _ms, _dtm, _tse, _dups) = per_shot_cache[d.name]
        for pool, series in ((dtheta_pool, (speed, heading)),
                             (dtheta_pool_raw,
                              cue_speed_series(track)[3:5])):
            sp, hd = series
            if len(sp) < 3:
                continue
            moving = sp > max(v_still_prop, V_INIT / 2)
            dth = np.abs(np.degrees(np.diff(np.unwrap(hd))))
            for i in range(len(dth)):
                if moving[i] and moving[i + 1]:
                    pool.append(dth[i])
    dtheta_pool = np.array(dtheta_pool) if dtheta_pool else np.array([0.0])
    dtheta_raw = (np.array(dtheta_pool_raw) if dtheta_pool_raw
                  else np.array([0.0]))
    kink_prop = float(np.percentile(dtheta_pool, 95))
    kink_prop_raw = float(np.percentile(dtheta_raw, 95))

    for d in shots:
        (track, dets, t_all, t, xy, t_mid, speed, heading, motion_start,
         dt_med, t_still_end, dups) = per_shot_cache[d.name]
        dur = float(t_all[-1] - t_all[0]) + dt_med
        covered = np.array([r["covered"] == "1" for r in track])
        dts = np.append(np.diff(t_all), dt_med)
        cov_time = float(dts[covered].sum() / dur)
        n_cue_frames = sum(1 for r in track if int(r["n_cands"]) > 0)
        gaps = [float(r["gap_s"]) for r in track if r["gap_s"]
                and float(r["gap_s"]) > 0]

        segs, kinks, dtheta, moving = ([], [], np.array([]), np.array([]))
        if len(speed) >= 3:
            segs, kinks, dtheta, moving = split_segments(
                t, xy, heading, speed, t_mid, v_still_prop, kink_prop)
        for a, b in segs:
            row = {"shot": d.name, "t_start_s": round(float(t[a]), 3)}
            row.update(fit_segment(t[a:b], xy[a:b]))
            seg_rows.append(row)

        sid = d.name.split("_")[1].lstrip("0") or "0"
        note = notes.get(sid, notes.get(d.name, {}))
        shot_rows.append({
            "shot": d.name, "n_frames": len(track),
            "dur_s": round(dur, 3),
            "cue_det_rate": round(n_cue_frames / len(track), 4),
            "coverage_time": round(cov_time, 4),
            "max_gap_s": round(max(gaps), 3) if gaps else 0.0,
            "motion_start_s": (round(motion_start - t_all[0], 3)
                               if motion_start is not None else ""),
            "n_kinks": len(kinks),
            "n_segments": len(segs),
            "n_dup_frames": len(dups),
            "contact_after_first": note.get("contact_after_first", ""),
            "note": note.get("comment", ""),
        })

        # plot per shot: speed trace + path
        fig, axes = plt.subplots(1, 2, figsize=(13, 4.5))
        if len(speed):
            axes[0].plot(t_mid - t_all[0], speed, lw=1)
            for k in kinks:
                axes[0].axvline(float(t[k] - t_all[0]), color="r", lw=0.6,
                                alpha=0.6)
            if motion_start is not None:
                axes[0].axvline(motion_start - t_all[0], color="g", ls="--",
                                lw=1, label="motion start")
            axes[0].axhline(v_still_prop, color="orange", ls=":",
                            label=f"v_still p99 = {v_still_prop:.3f} m/s")
            axes[0].legend(fontsize=8)
        axes[0].set_xlabel("t (s)")
        axes[0].set_ylabel("cue speed (m/s)")
        axes[0].set_title(f"{d.name} speed trace (kinks do)")
        if len(xy):
            axes[1].plot(xy[:, 0], xy[:, 1], ".-", ms=2, lw=0.8)
            axes[1].plot(xy[0, 0], xy[0, 1], "go", label="start")
            axes[1].legend(fontsize=8)
        axes[1].set_xlim(-0.05, 1.32)
        axes[1].set_ylim(-0.05, 2.59)
        axes[1].set_aspect("equal")
        axes[1].set_title("cue path (table m)")
        fig.tight_layout()
        fig.savefig(plots / f"{d.name}_trace.png", dpi=110)
        plt.close(fig)

    # ---- histogram tong hop --------------------------------------------
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    if len(still_sp) > 1:
        axes[0].hist(still_sp * 1000, bins=40)
    axes[0].axvline(v_still_prop * 1000, color="r",
                    label=f"p99 = {v_still_prop * 1000:.1f} mm/s")
    axes[0].set_xlabel("apparent speed khi DUNG YEN (mm/s)")
    axes[0].legend(fontsize=8)
    if len(dtheta_pool) > 1:
        axes[1].hist(dtheta_pool, bins=40)
    axes[1].axvline(kink_prop, color="r", label=f"p95 = {kink_prop:.1f} deg")
    axes[1].set_xlabel("|delta heading| giua 2 buoc khi DANG CHAY (deg, da loai frame trung)")
    axes[1].legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(plots / "thresholds_hist.png", dpi=110)
    plt.close(fig)

    # ---- summary + G0 ---------------------------------------------------
    write_csv(args.out_root / "bb9_p0_shots.csv", shot_rows)
    write_csv(args.out_root / "bb9_p0_noise.csv", noise_rows)
    write_csv(args.out_root / "bb9_p0_segments.csv", seg_rows)

    n = len(shot_rows)
    n80 = sum(1 for r in shot_rows if r["coverage_time"] >= 0.80)
    dev_pool = np.concatenate(all_dev) if all_dev else np.array([])
    contact_vals = [r["contact_after_first"] for r in shot_rows
                    if r["contact_after_first"] != ""]
    n_noc = sum(1 for v in contact_vals if v is False or v == "False")

    n_dup_total = sum(r["n_dup_frames"] for r in shot_rows)
    n_frames_total = sum(r["n_frames"] for r in shot_rows)
    summary = {"n_shots": n, "n_cov80": n80,
               "G0_track": f"{n80}/{n} >= 7/10: "
                           f"{'DAT' if n80 >= 7 and n >= 10 else 'DO'}",
               "v_still_prop_mms": round(v_still_prop * 1000, 1),
               "kink_prop_deg": round(kink_prop, 1),
               "kink_prop_raw_deg": round(kink_prop_raw, 1),
               "dup_frame_rate": round(n_dup_total / max(n_frames_total, 1), 4),
               "n_noise_balls": len(noise_rows),
               "n_segments": len(seg_rows),
               "spin_unident_noted": (f"{n_noc}/{len(contact_vals)}"
                                      if contact_vals else "chua ghi chu")}
    if len(dev_pool):
        summary.update(pct_dict(dev_pool, "noise_r_mm", scale=1000.0))
    write_csv(args.out_root / "bb9_p0_summary.csv", [summary])

    print("\n===== G0 =====")
    for k, v in summary.items():
        print(f"  {k}: {v}")
    print(f"\nCSV: {args.out_root}\\bb9_p0_*.csv   plots: {plots}")


if __name__ == "__main__":
    main()