src/poolcoach_cv/resim.py

"""Resim + RMSE — lát A2 phần 2 (13/08/2026, BRIEF A2).

Sau khi analyze xong MỘT cú, nạp params suy được vào pooltool sim lại cú
trên bàn kích thước broadcast (2.54 × 1.27 m — cùng hệ toạ độ track) rồi
đo RMSE (mm, theo điểm track cue ball) giữa quỹ đạo resim và quan sát.
Chạy trong CV WORKER (venv ``poolcoach-cv-env`` — pooltool cài thêm ở
requirements-cv.txt, lát A2): worker là nơi duy nhất còn giữ đủ ngữ cảnh
per cú trước khi clip bị xoá, và JSON per cú phải ra ĐĨA ngay trong job
(nếp BG29b/BG31 — artifact ghi file ngay).

HAI bộ params, KHÔNG trộn (BRIEF A2 bối cảnh 3):

- ``shotnet``  — V0 GẬY (đã quy đổi m/s bàn thật), φ, (a, b) từ model.
- ``analytic`` — V0 là tốc độ BI đo từ track (thước KHÁC): quy về V0 gậy
  bằng tỉ lệ ``v_ball/V0`` đo một lần lúc warmup (strike giữa bi trên
  chính bàn resim — thuần pooltool, không fit); spin GIẢ ĐỊNH a=b=0
  (analytic không có số a/b — chỉ có lớp spin thô). Hai giả định này ghi
  thẳng vào JSON (``assumptions``) — người đọc RMSE analytic phải biết nó
  gánh thêm mô hình hoá, không so ngang mù với RMSE shotnet.

Neo thời gian: sim t=0 là lúc gậy chạm bi; quan sát chỉ có ``motion_start``
(phát hiện SAU cú chạm tới ~1 bước frame — V_INIT_RUN bước vượt ngưỡng).
Nên t0 được DÒ trên lưới [0, 2·dt] trước motion_start (ALIGN_*), chọn nấc
RMSE nhỏ nhất, và ghi ``t_align_s`` vào JSON — khớp là khớp HÌNH quỹ đạo,
không để lệch pha lấy mẫu ~30–70mm giả dạng sai số suy ngược. Lưới dò là
tất định và được khai — không phải "vá để lấy số đẹp".

Series resim (toạ độ bàn + t, mọi bi dịch chuyển đủ xa) lưu LUÔN vào JSON
— trang chi tiết (nét đứt) và nút "Đánh lại cú" dùng chung, không tính
hai lần (BRIEF A2 phần 2).

Import pooltool LƯỜI trong hàm — venv app import module này được (tests
contract chạy không pooltool thật, nếp poolcoach_cv); phần toán thuần
(``rmse_mm``, ``align_t0``, ``sample_traj``) test được bằng quỹ đạo giả.
"""

from __future__ import annotations

import math
import time

import numpy as np

# Bàn resim = bàn broadcast (poolcoach_cv.broadcast TABLE_*) — nhắc lại số
# ở đây sẽ thành bản sao có thể trôi, nên import từ MỘT nguồn.
from poolcoach_cv.broadcast import BALL_R_M, TABLE_L_M, TABLE_W_M

RESIM_DT_S = 0.02        # bước lấy mẫu series ghi JSON (50 Hz — đủ mượt để
                         # vẽ nét đứt + animate, JSON không phình)
RESIM_MIN_DISP_M = 0.02  # bi resim dịch tổng dưới mức này không ghi series
                         # (bi đứng yên — series chỉ tổ phình JSON)
ALIGN_STEPS = 8          # dò t0 trên lưới ALIGN_STEPS+1 nấc đều...
ALIGN_SPAN_DT = 2.0      # ...phủ [0, ALIGN_SPAN_DT·dt_med] TRƯỚC motion_start
                         # (motion_start trễ sau cú chạm ≤ ~2 bước frame —
                         # V_INIT_RUN=3 bước, neo lấy bước đầu của run)
WARMUP_V0 = 2.0          # strike đo tỉ lệ v_ball/V0 (giữa bi, bàn resim)


# ------------------------------------------------------- toán thuần (no pt)

def sample_traj(sim_t: np.ndarray, sim_xy: np.ndarray,
                at_t: np.ndarray) -> np.ndarray:
    """Nội suy tuyến tính quỹ đạo sim tại các mốc ``at_t`` (giây, trục sim).
    Ngoài hai đầu → kẹp về điểm đầu/cuối (bi chưa đánh / đã dừng): so với
    quan sát vẫn ĐÚNG nghĩa — quan sát còn chạy mà sim đã dừng thì lệch đó
    là lệch thật, không được cắt bỏ."""
    return np.stack([np.interp(at_t, sim_t, sim_xy[:, 0]),
                     np.interp(at_t, sim_t, sim_xy[:, 1])], axis=1)


def rmse_mm(obs_xy: np.ndarray, sim_xy: np.ndarray) -> float:
    """RMSE (mm) theo điểm giữa hai chuỗi vị trí cùng độ dài."""
    d = np.asarray(obs_xy, dtype=np.float64) - np.asarray(sim_xy,
                                                          dtype=np.float64)
    return float(np.sqrt(np.mean(d[:, 0] ** 2 + d[:, 1] ** 2)) * 1000.0)


def align_t0(obs_t: np.ndarray, obs_xy: np.ndarray, sim_t: np.ndarray,
             sim_xy: np.ndarray, motion_start: float,
             dt_med: float) -> tuple[float, float, float]:
    """Dò mốc t0 của sim trên trục track: t0 = motion_start − Δ, Δ trên
    lưới đều [0, ALIGN_SPAN_DT·dt_med] (ALIGN_STEPS+1 nấc). Trả
    ``(t0, t_align=Δ, rmse_mm)`` của nấc RMSE nhỏ nhất — tất định, khai
    trong JSON."""
    best = None
    for k in range(ALIGN_STEPS + 1):
        delta = ALIGN_SPAN_DT * dt_med * k / ALIGN_STEPS
        t0 = motion_start - delta
        err = rmse_mm(obs_xy, sample_traj(sim_t, sim_xy, obs_t - t0))
        if best is None or err < best[2]:
            best = (t0, delta, err)
    return best


def _series(sim_t: np.ndarray, sim_xy: np.ndarray, t0: float) -> list:
    """Series một bi cho JSON: [[t_s, x_m, y_m], ...] lấy mẫu RESIM_DT_S,
    t_s trên CÙNG trục thời gian với ``track`` của kết quả analyze."""
    ts = np.arange(0.0, float(sim_t[-1]) + RESIM_DT_S / 2, RESIM_DT_S)
    xy = sample_traj(sim_t, sim_xy, ts)
    return [[round(float(t0 + t), 3), round(float(x), 4), round(float(y), 4)]
            for t, (x, y) in zip(ts, xy)]


# --------------------------------------------------- pooltool (import lười)

def _build_system(cue_xy, others_xy, v0, phi, a, b, state):
    import pooltool as pt

    balls = {"cue": pt.Ball.create("cue", xy=(float(cue_xy[0]),
                                              float(cue_xy[1])))}
    for i, xy in enumerate(others_xy, start=1):
        balls[f"b{i}"] = pt.Ball.create(f"b{i}", xy=(float(xy[0]),
                                                     float(xy[1])))
    system = pt.System(table=state["table"], balls=balls,
                       cue=pt.Cue(cue_ball_id="cue"))
    system.cue.set_state(V0=float(v0), phi=float(phi) % 360.0,
                         a=float(a), b=float(b))
    return system


def _sim_trajs(cue_xy, others_xy, v0, phi, a, b, state) -> dict:
    """Sim MỘT cú → {ball_id: (ts, xy (n,2))} theo history continuize.
    Ném exception khi pooltool ném — caller khai ``error`` per bộ, worker
    không được chết vì một bộ params rác (cú ma vẫn được resim thử)."""
    import pooltool as pt

    system = _build_system(cue_xy, others_xy, v0, phi, a, b, state)
    pt.simulate(system, inplace=True)
    pt.continuize(system, dt=0.01, inplace=True)
    out = {}
    for bid, ball in system.balls.items():
        rvw, _ss, ts = ball.history_cts.vectorize()
        out[bid] = (np.asarray(ts, dtype=np.float64),
                    np.asarray(rvw[:, 0, :2], dtype=np.float64))
    return out


def warmup() -> dict | None:
    """Dựng bàn resim + JIT pooltool MỘT lần lúc boot worker (nếp engine
    worker: không báo alive lúc còn đang JIT — sim nguội 40–90s, sim ấm
    ~40ms). Warmup phủ các resolver hay gặp: chạm bi, chạm băng, spin,
    vào lỗ. Đo luôn tỉ lệ ``v_ball/V0`` (strike giữa bi) cho bộ analytic.
    Trả state dict hoặc None khi pooltool vắng/hỏng (resim tắt, analyze
    sống nguyên — nếp shotnet)."""
    try:
        import pooltool as pt
        from pooltool.objects.table.specs import PocketTableSpecs
    except Exception as e:  # noqa: BLE001 — thiếu pooltool ≠ worker chết
        print(f"[resim] pooltool khong import duoc ({type(e).__name__}: {e})"
              f" -- resim TAT, analyze van chay", flush=True)
        return None
    t0 = time.time()
    state = {"table": pt.Table.from_table_specs(
        PocketTableSpecs(l=TABLE_L_M, w=TABLE_W_M))}
    try:
        # 1) strike giữa bi, không chạm gì — đo tỉ lệ v_ball/V0
        trajs = _sim_trajs((TABLE_W_M / 2, 0.8), [], WARMUP_V0, 90.0,
                           0.0, 0.0, state)
        ts, xy = trajs["cue"]
        v = np.hypot(*np.diff(xy, axis=0).T) / np.maximum(np.diff(ts), 1e-9)
        state["v_ball_per_v0"] = round(float(v[0]) / WARMUP_V0, 4)
        # 2) chạm bi + spin + băng + dồn về góc lỗ — JIT các resolver còn lại
        _sim_trajs((TABLE_W_M / 2, 0.6), [(TABLE_W_M / 2 + 0.02, 1.4)],
                   3.0, 90.0, 0.25, 0.3, state)
        _sim_trajs((0.3, 0.4), [(1.0, 2.2)], 4.0, 60.0, -0.3, -0.3, state)
    except Exception as e:  # noqa: BLE001
        print(f"[resim] warmup pooltool hong ({type(e).__name__}: {e}) "
              f"-- resim TAT", flush=True)
        return None
    print(f"[resim] pooltool JIT xong sau {time.time() - t0:.1f}s "
          f"(v_ball/V0 = {state['v_ball_per_v0']:g}, ban "
          f"{TABLE_W_M}x{TABLE_L_M}m)", flush=True)
    return state


# ------------------------------------------------------------- khối resim

def _obs_arrays(result: dict):
    """(obs_t, obs_xy) các điểm track TỪ motion_start trở đi + vị trí cue
    lúc nghỉ (median track trước motion_start; thiếu thì điểm đầu cửa sổ)."""
    ms = (result.get("metrics") or {}).get("motion_start_s")
    track = result.get("track") or []
    t = np.array([float(p["t_s"]) for p in track])
    xy = np.array([[float(p["x_m"]), float(p["y_m"])] for p in track]
                  ) if track else np.empty((0, 2))
    after = t >= ms - 1e-9
    before = ~after
    if int(before.sum()) >= 1:
        cue0 = np.median(xy[before], axis=0)
    elif int(after.sum()):
        cue0 = xy[after][0]
    else:
        cue0 = None
    return t[after], xy[after], cue0


def resim_result(result: dict, state: dict) -> dict:
    """Khối ``resim`` cho JSON kết quả MỘT cú — hai bộ params độc lập,
    mỗi bộ hoặc đủ (params/rmse/series) hoặc ``error`` nói vì sao vắng.
    Không đụng một số nào có sẵn trong ``result``."""
    m = result.get("metrics") or {}
    out: dict = {"table_w_m": TABLE_W_M, "table_l_m": TABLE_L_M,
                 "ball_r_m": BALL_R_M, "sets": {}}

    ms = m.get("motion_start_s")
    if ms is None:
        why = ("không có motion_start để neo t=0 của sim — cú không resim "
               "được (chính nó cũng là tín hiệu van)")
        out["sets"]["shotnet"] = {"error": why}
        out["sets"]["analytic"] = {"error": why}
        return out

    obs_t, obs_xy, cue0 = _obs_arrays(result)
    if len(obs_t) < 2 or cue0 is None:
        why = "không đủ điểm track sau motion_start để so quỹ đạo"
        out["sets"]["shotnet"] = {"error": why}
        out["sets"]["analytic"] = {"error": why}
        return out

    dt_med = (float(np.median(np.diff(obs_t))) if len(obs_t) > 1
              else 1 / 30.0)
    others = [(float(b["x_m"]), float(b["y_m"]))
              for b in result.get("balls_init") or []
              if b.get("type") != "cue"
              and math.hypot(float(b["x_m"]) - cue0[0],
                             float(b["y_m"]) - cue0[1]) > 2 * BALL_R_M]

    sets: dict[str, dict] = {}
    sn = result.get("shotnet")
    if sn:
        sets["shotnet"] = {
            "params": {"v0_mps": float(sn["v0_cue_mps"]),
                       "phi_deg": float(sn["phi_deg"]),
                       "a": float(sn["a"]), "b": float(sn["b"])},
            "assumptions": "V0 gậy (model, đã quy đổi m/s bàn thật); "
                           "a/b từ model",
        }
    else:
        sets["shotnet"] = {"error": "kết quả không có khối shotnet"}
    v0_ball = m.get("v0_mps")
    phi_an = m.get("phi_deg")
    if v0_ball is not None and phi_an is not None:
        sets["analytic"] = {
            "params": {"v0_mps": round(float(v0_ball)
                                       / state["v_ball_per_v0"], 3),
                       "phi_deg": float(phi_an), "a": 0.0, "b": 0.0},
            "assumptions": f"a=b=0 (analytic không có số a/b); V0 gậy suy "
                           f"từ tốc độ bi đo qua tỉ lệ v_ball/V0 = "
                           f"{state['v_ball_per_v0']:g} (đo lúc warmup, "
                           f"strike giữa bi)",
        }
    else:
        sets["analytic"] = {"error": "analytic không đo được V0/φ"}

    for name, blk in sets.items():
        if "params" not in blk:
            continue
        p = blk["params"]
        try:
            trajs = _sim_trajs(cue0, others, p["v0_mps"], p["phi_deg"],
                               p["a"], p["b"], state)
        except Exception as e:  # noqa: BLE001 — params rác ≠ worker chết
            blk.pop("params", None)
            blk["error"] = (f"pooltool không sim được bộ {name} "
                            f"({type(e).__name__}: {e})")
            continue
        sim_t, sim_xy = trajs["cue"]
        t0, dl, err = align_t0(obs_t, obs_xy, sim_t, sim_xy, float(ms),
                               dt_med)
        blk["rmse_mm"] = round(err, 1)
        blk["n_points"] = int(len(obs_t))
        blk["t0_s"] = round(t0, 3)
        blk["t_align_s"] = round(dl, 3)
        series = {}
        for bid, (bt, bxy) in trajs.items():
            disp = float(np.linalg.norm(np.diff(bxy, axis=0),
                                        axis=1).sum()) if len(bxy) > 1 else 0
            if bid == "cue" or disp >= RESIM_MIN_DISP_M:
                series[bid] = _series(bt, bxy, t0)
        blk["series"] = series
    out["sets"] = sets
    return out