examples/pooltool_demo.py

#!/usr/bin/env python3
"""
Demo pooltool cho PoolCoach — minh hoạ những hook cần cho RL position play.

Cài đặt (máy local, nên dùng Python 3.10-3.12):
    pip install pooltool-billiards

Chạy headless (không GUI):
    python pooltool_demo.py

Bật GUI 3D xem cú đánh (cần màn hình; ESC để thoát):
    python pooltool_demo.py --gui
"""

import sys

import pooltool as pt


def build_nineball_shot():
    """Dựng một hệ 9-ball chuẩn + một cú đánh có spin."""
    table = pt.Table.default()
    balls = pt.get_rack(pt.GameType.NINEBALL, table)   # rack 9-ball đúng luật
    cue = pt.Cue(cue_ball_id="cue")
    shot = pt.System(table=table, balls=balls, cue=cue)

    # === ĐÂY LÀ "ACTION" CỦA AGENT ===
    # V0 = lực, phi = góc hướng, a = english (side spin), b = draw/follow
    shot.cue.set_state(
        V0=8.0,                          # power
        phi=pt.aim.at_ball(shot, "1"),   # angle: nhắm bi số 1 (luật 9-ball: bi nhỏ nhất trước)
        a=0.0,                           # spin_x: english trái/phải
        b=0.2,                           # spin_y: follow (dương) / draw (âm) -> điều bi cái
    )
    return shot


def observation(shot):
    """OBSERVATION cho RL: toạ độ (x, y) của mọi bi còn trên bàn."""
    obs = {}
    for ball_id, ball in shot.balls.items():
        xyz = ball.state.rvw[0]          # vị trí [x, y, z]
        obs[ball_id] = (round(float(xyz[0]), 4), round(float(xyz[1]), 4))
    return obs


def summarize_events(shot):
    """REWARD source: đọc events sau cú đánh."""
    pocketed, cushions, ball_hits = [], 0, 0
    for e in shot.events:
        etype = str(e.event_type)
        if "pocket" in etype.lower():
            pocketed.append([a.id for a in e.agents])
        elif "cushion" in etype.lower():
            cushions += 1
        elif etype.lower().endswith("ball_ball"):
            ball_hits += 1
    return pocketed, cushions, ball_hits


def main():
    shot = build_nineball_shot()

    print("== Vị trí bi TRƯỚC cú đánh ==")
    for bid, xy in observation(shot).items():
        print(f"  {bid:>4}: {xy}")

    # === "STEP": mô phỏng headless, rất nhanh nhờ event-based + Numba ===
    pt.simulate(shot, inplace=True)

    print(f"\nSố sự kiện mô phỏng: {len(shot.events)}")
    print(f"Thời gian mô phỏng (giây): {shot.t:.3f}")

    pocketed, cushions, ball_hits = summarize_events(shot)
    print("\n== Tín hiệu cho REWARD ==")
    print(f"  Bi vào lỗ: {pocketed if pocketed else 'không'}")
    print(f"  Số lần chạm băng: {cushions}")
    print(f"  Số va chạm bi-bi: {ball_hits}")

    print("\n== Vị trí bi SAU cú đánh (đặc biệt xem 'cue' -> position play) ==")
    for bid, xy in observation(shot).items():
        print(f"  {bid:>4}: {xy}")

    # GUI 3D: KHÔNG simulate lại (shot đã simulate ở trên — hệ hết năng
    # lượng, simulate nữa là rỗng); pt.show tự continuize theo fps.
    if "--gui" in sys.argv:
        pt.show(shot, title="PoolCoach demo — break 9-ball")


if __name__ == "__main__":
    main()