src/poolcoach_rl/recommend/simulate.py

"""Tầng SIM — chỉ nói sự thật vật lý, không luật, không xếp hạng.

24/07/2026 (FullRack phase A): `simulate_shot_multi` nhận dict {id: xy} bất
kỳ, trả FACTS (potted list / scratch / first_contact / balls_final).
31/07/2026 thêm hai facts cho V2 (`bb_events`, `cue_roll_after` — xem
docstring hàm). `simulate_shot` 3-bi ở lại làm wrapper cho `warmup`.

31/07/2026 (việc D): các khoá xếp hạng EV (`rank_key`, `rank_key_factory`,
`rank2_key`, `EV2_MAX`) xoá theo engine oracle/hybrid — chúng sống ở nhánh
`v1-full`. POS_COEF/SCRATCH_EV Ở LẠI: `rules.judge_shot` và reward env vẫn
định nghĩa Q/EV bằng chúng (hồ sơ nghiên cứu, xem BRIEF bối cảnh 5).
"""

from __future__ import annotations

POS_COEF = 0.5          # EV = pot·(1 + POS_COEF·Q) — khớp reward env
SCRATCH_EV = -0.3       # khớp scratch_penalty env


def path_len_after(xy, ts, t1, *, eps=1e-12):
    """Độ dài polyline (mét) của ``xy`` từ thời điểm ``t1`` trở đi — THUẦN.

    ``xy``: mảng ``(n, 2)`` vị trí theo thời gian; ``ts``: ``(n,)`` mốc thời
    gian tương ứng. Nguồn dự kiến là ``ball.history`` event-based của pooltool
    (KHÔNG cần continuize): snapshot tại MỌI event nên mọi cú dội băng và mọi
    chuyển pha sliding→rolling đều là đỉnh polyline. Nghiệm thu 31/07/2026 so
    với continuize ``dt=0.01`` trên 6 cú trộn xoáy/lực: lệch 0.2–8.9 mm —
    nhỏ hơn nhiều bước lưới ứng viên, đủ cho xếp hạng.

    ``eps`` nới mốc so sánh: ``t1`` chép từ ``event.time`` và ``ts`` chép từ
    cùng nguồn, nhưng qua hai đường float khác nhau — so ``>=`` trần trụi có
    thể rớt đúng snapshot va chạm.
    """
    import numpy as np

    xy = np.asarray(xy, dtype=np.float64).reshape(-1, 2)
    m = np.asarray(ts, dtype=np.float64).ravel() >= float(t1) - eps
    if int(m.sum()) < 2:
        return 0.0
    return float(np.linalg.norm(np.diff(xy[m], axis=0), axis=1).sum())


def simulate_shot_multi(env_h, balls, phi, v0, a, b, render=False):
    """1 cú với N bi bất kỳ. ``balls = {"cue": xy, "1": xy, "5": xy, ...}``
    (bắt buộc có "cue"; id bi = chuỗi số như FE/API gửi).

    Chỉ trả FACTS vật lý — không luật, không EV (tầng đó ở phase B):
      system        : pooltool System nếu render=True (vẽ quỹ đạo) | None
      potted        : list id bi VÀO LỖ trong cú này (không tính cue)
      scratch       : bool — bi cái vào lỗ
      first_contact : id bi đầu tiên cue chạm | None nếu không chạm gì
                      (event BALL_BALL đầu tiên có "cue" — luật 9-ball
                      phase B sẽ so với target để bắt foul)
      balls_final   : {id: np.ndarray (2,) | None nếu vào lỗ} — mọi bi,
                      kể cả cue (nút "Đánh cú này" của FE dùng dict này)

    HAI KEY THÊM 31/07/2026 (ZonePlanner V2, BRIEF bối cảnh 7 — MỞ RỘNG
    THUẦN: hai phép gán, không nhánh cũ nào đọc chúng):
      bb_events     : tuple[(t, (id, id))] — MỌI event BALL_BALL theo thứ tự
                      thời gian. Lọc (i)/(iii) của V2 đọc từ đây ("đường bi
                      mục tiêu không chạm bi khác", "bi cái sau va chạm không
                      đụng bi nào — kể cả double-kiss").
      cue_roll_after: float | None — quãng đường bi cái LĂN SAU cú chạm bi
                      đầu tiên (mét, arc length event-polyline, xem
                      ``path_len_after``) | None nếu cue không chạm bi nào.
                      Khoá xếp hạng số 1 của V2 (§8.2 design).

    Trả None nếu pooltool ném exception (giữ hợp đồng cũ).
    """
    import numpy as np
    import pooltool as pt
    import pooltool.constants as ptc

    system = pt.System(
        table=env_h.table,
        balls={bid: pt.Ball.create(bid, xy=tuple(xy))
               for bid, xy in balls.items()},
        cue=pt.Cue(cue_ball_id="cue"),
    )
    system.cue.set_state(V0=v0, phi=phi, a=a, b=b)
    try:
        pt.simulate(system, inplace=True)
    except Exception:
        return None
    if render:
        try:
            pt.continuize(system, dt=0.01, inplace=True)
        except Exception:
            pass  # fallback: vẽ theo history event-based

    def _pk(bid):
        return system.balls[bid].state.s == ptc.pocketed

    def _xy(bid):
        return np.asarray(system.balls[bid].state.rvw[0][:2], dtype=np.float64)

    first_contact = None
    for ev in system.events:
        if ev.event_type == pt.EventType.BALL_BALL and "cue" in ev.ids:
            first_contact = next(i for i in ev.ids if i != "cue")
            break

    bb_events = tuple((float(ev.time), tuple(ev.ids))
                      for ev in system.events
                      if ev.event_type == pt.EventType.BALL_BALL)

    # quãng đường bi cái SAU cú chạm đầu — từ history event-based (có sẵn
    # không cần continuize; extract_trajectories vẫn fallback vào đúng nguồn
    # này). vectorize hỏng thì trả None thay vì làm chết cả sim: hai key mới
    # là mở rộng, không được phép kéo sập hợp đồng cũ.
    cue_roll_after = None
    t1 = next((t for t, ids in bb_events if "cue" in ids), None)
    if t1 is not None:
        try:
            rvw, _ss, ts = system.balls["cue"].history.vectorize()
        except Exception:
            rvw = None
        if rvw is not None and len(rvw) > 1:
            cue_roll_after = path_len_after(rvw[:, 0, :2], ts, t1)

    return {
        "system": system if render else None,
        "potted": [bid for bid in balls if bid != "cue" and _pk(bid)],
        "scratch": _pk("cue"),
        "first_contact": first_contact,
        "balls_final": {bid: (None if _pk(bid) else _xy(bid))
                        for bid in balls},
        "bb_events": bb_events,
        "cue_roll_after": cue_roll_after,
    }


def simulate_shot(env_h, cue_xy, b1_xy, b2_xy, phi, v0, a, b, render=False):
    """1 cú 3 bi (wrapper tương thích v1) — output y hệt bản cũ 23/07,
    thêm 2 key mới vô hại: ``first_contact``, ``balls_final``."""
    balls = {"cue": cue_xy, "1": b1_xy}
    if b2_xy is not None:
        balls["2"] = b2_xy
    m = simulate_shot_multi(env_h, balls, phi, v0, a, b, render=render)
    if m is None:
        return None

    potted = "1" in m["potted"]
    scratch = m["scratch"]
    b2p = "2" in m["potted"]
    bf = m["balls_final"]
    q = 0.0
    if b2_xy is not None and potted and not scratch:
        q = 1.0 if b2p else env_h._position_q(bf["cue"], bf["2"])
    if scratch:
        ev = SCRATCH_EV
    elif potted:
        ev = 1.0 + POS_COEF * q
    else:
        ev = 0.0
    return {
        "system": m["system"],
        "potted": potted, "scratch": scratch, "b2_potted": b2p,
        "q": q, "ev": ev,
        "cue_final": bf["cue"],  # None khi scratch (bi cái pocketed) — như cũ
        "b2_final": bf["2"] if b2_xy is not None else None,  # None khi b2 vào lỗ
        "first_contact": m["first_contact"],
        "balls_final": bf,
    }