"""Benchmark tốc độ sim N bi — việc ĐẦU TIÊN của FullRack phase A (§4.4).

Rủi ro chính của full rack: sim 10 bi nhiều event hơn hẳn 3 bi. Đo ms/cú
trên máy thật (sandbox không có pooltool) cho 3 kịch bản:

  1. ref_3bi     — cue+1+2 random (mốc so sánh với app hiện tại)
  2. midgame_10  — 10 bi rải random, đánh vào bi 1 (kịch bản search chính)
  3. break_rack  — cue đánh vỡ rack 9 bi nguyên (worst case nhiều event)

In kèm ước lượng thời gian search oracle/cú (grid 250 cell × số lỗ khả thi,
giả định 3 lỗ = 750 sim) — ngưỡng chấp nhận theo design: < ~5s/cú.

Chạy:  python scripts/benchmark_fullrack.py --shots 200
"""

from __future__ import annotations

import argparse
import math
import os
import statistics
import sys
import time
from pathlib import Path

import numpy as np

sys.stdout.reconfigure(encoding="utf-8")
sys.stderr.reconfigure(encoding="utf-8")

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

SEARCH_SIMS = 750  # 3 lỗ khả thi × 250 cell (10 V0 × 5 side × 5 vert)


def _place_random(rng, w, l, n, ball_r, min_gap=1.1):
    """n vị trí không chồng nhau (tâm cách nhau ≥ 2R×min_gap, né sát mép)."""
    pts = []
    while len(pts) < n:
        p = np.array([rng.uniform(2 * ball_r, w - 2 * ball_r),
                      rng.uniform(2 * ball_r, l - 2 * ball_r)])
        if all(np.linalg.norm(p - q) >= 2 * ball_r * min_gap for q in pts):
            pts.append(p)
    return pts


def _aim_at(balls, target, rng, noise_deg):
    d = balls[target] - balls["cue"]
    return (math.degrees(math.atan2(d[1], d[0]))
            + rng.uniform(-noise_deg, noise_deg)) % 360.0


def _bench(env_h, name, gen_case, n_shots, rng):
    from poolcoach_rl.recommend.simulate import simulate_shot_multi

    times, fails, contacts = [], 0, 0
    for _ in range(n_shots):
        balls, phi, v0, a, b = gen_case(rng)
        t0 = time.perf_counter()
        r = simulate_shot_multi(env_h, balls, phi, v0, a, b)
        times.append((time.perf_counter() - t0) * 1000.0)
        if r is None:
            fails += 1
        elif r["first_contact"] is not None:
            contacts += 1

    med = statistics.median(times)
    p90 = statistics.quantiles(times, n=10)[-1]
    sps = 1000.0 / med
    ncpu = os.cpu_count() or 1
    est1 = SEARCH_SIMS * med / 1000.0
    print(f"=== {name} ({n_shots} cú) ===")
    print(f"  mean {statistics.mean(times):7.1f} ms | median {med:7.1f}"
          f" | p90 {p90:7.1f} | max {max(times):7.1f}"
          f"  ->  {sps:6.1f} sim/s (1 core)")
    print(f"  contact {contacts}/{n_shots}, sim fail {fails}")
    print(f"  search ~{SEARCH_SIMS} sim: {est1:6.1f} s (1 core)"
          f" | ~{est1 / ncpu:5.1f} s ({ncpu} core lý thuyết)")
    print()
    return med


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--shots", type=int, default=200, help="số cú/kịch bản")
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    import pooltool as pt

    from poolcoach_rl.envs import PositionPlayEnv
    from poolcoach_rl.envs.position_env import BALL_R
    from poolcoach_rl.recommend.core import warmup

    env_h = PositionPlayEnv()
    w, l = env_h.w, env_h.l
    rng = np.random.default_rng(args.seed)

    print("JIT warmup (~40s lần đầu)...", flush=True)
    print(f"  xong sau {warmup(env_h):.1f}s\n")

    # Rack 9 bi chuẩn từ pooltool; nếu get_rack không kèm cue thì đặt
    # cue phía đối diện rack (chỉ cần hợp lý cho benchmark).
    rack = pt.get_rack(pt.GameType.NINEBALL, env_h.table)
    rack_xy = {bid: np.asarray(bb.state.rvw[0][:2], dtype=np.float64)
               for bid, bb in rack.items()}
    if "cue" not in rack_xy:
        cue_y = l - rack_xy["1"][1]
        rack_xy["cue"] = np.array(
            [w / 2, min(max(cue_y, 2 * BALL_R), l - 2 * BALL_R)])
    n_rack = len(rack_xy)
    print(f"rack: {n_rack} bi ({', '.join(sorted(rack_xy))})\n")

    ids10 = ["cue"] + [str(i) for i in range(1, 10)]

    def case_ref3(rng):
        balls = dict(zip(["cue", "1", "2"],
                         _place_random(rng, w, l, 3, BALL_R)))
        return (balls, _aim_at(balls, "1", rng, 30.0),
                rng.uniform(0.5, 4.0),
                rng.uniform(-0.4, 0.4), rng.uniform(-0.4, 0.4))

    def case_mid10(rng):
        balls = dict(zip(ids10, _place_random(rng, w, l, 10, BALL_R)))
        return (balls, _aim_at(balls, "1", rng, 30.0),
                rng.uniform(0.5, 4.0),
                rng.uniform(-0.4, 0.4), rng.uniform(-0.4, 0.4))

    def case_break(rng):
        balls = {k: v.copy() for k, v in rack_xy.items()}
        return (balls, _aim_at(balls, "1", rng, 5.0),
                rng.uniform(2.5, 4.0),
                rng.uniform(-0.2, 0.2), rng.uniform(-0.2, 0.2))

    m3 = _bench(env_h, "ref_3bi (mốc app hiện tại)", case_ref3,
                args.shots, rng)
    m10 = _bench(env_h, "midgame_10bi (kịch bản search chính)", case_mid10,
                 args.shots, rng)
    mbr = _bench(env_h, f"break_rack ({n_rack} bi, worst case)", case_break,
                 max(args.shots // 4, 20), rng)

    print("=== KẾT LUẬN ===")
    print(f"  10 bi chậm hơn 3 bi: ×{m10 / m3:.1f} (midgame), "
          f"×{mbr / m3:.1f} (break)")
    verdict = ("ĐẠT — search midgame 1 core đã < 5s/cú"
               if SEARCH_SIMS * m10 / 1000.0 < 5.0 else
               "CHƯA ĐẠT ngưỡng 5s/cú 1 core → cần ProcessPool trong BE "
               "(mitigation 1, §4.4) hoặc coarse-to-fine (mitigation 2)")
    print(f"  {verdict}")


if __name__ == "__main__":
    main()