# -*- coding: utf-8 -*-
"""Sinh dataset synthetic BB9 cho net suy ngược P2 (BG26, 11/08/2026).

Pipeline mỗi cú:  thế bàn random (tái dụng ``_place_random`` của
``scripts/eval_fullrack.py`` — cùng generator đã nghiệm thu engine) + action
4D đúng dải env RL → pooltool simulate (bàn default, KHÔNG chỉnh physics —
BRIEF bối cảnh 6) → quỹ đạo lý tưởng continuize 100Hz → LÀM BẨN theo mô hình
nhiễu P0 (bảng hằng NOISE_* dưới, số gốc HANDOFF 23) → shard ``.npz``.

Nhiễu P0 là HỢP ĐỒNG (BRIEF bối cảnh 3) — ba đặc trưng bắt buộc, thiếu một
là net BG27 đẹp trên sim chết trên thật (bài Kienzle 2023):
  1. ellipse DỊ HƯỚNG xoay theo hướng vận tốc (dọc ~±13mm @2 m/s scale tuyến
     tính theo tốc độ, ngang ~±0.8mm — HANDOFF 23 Bất ngờ 2);
  2. chu kỳ frame TRÙNG 1/6 kiểu 25→30 upconvert khi fps=30 (dup-rate thô
     16.7%, đo được qua img_diff ~11% — Bất ngờ 1): frame trùng CHÉP NGUYÊN
     vị trí đã nhiễu của frame trước + img_diff nhỏ, PTS vẫn đều;
  3. dropout ĐẦU cú 0–25% thời lượng chuyển động của cue (người cúi ngắm che
     — shot_07 mất 24% đầu cú, Bất ngờ 3) + gap ngắn ≤0.3s rải ngẫu nhiên.

Label: (V0, phi, a, b) = ACTION đưa vào ``cue.set_state`` — V0 là tốc độ
GẬY. Probe strike (BG26b, ghi trong spec): tốc độ BI rời gậy =
V0 · 2/(1 + m/M + 2.5·(a²+b²)) = V0 · 2/(1.3 + 2.5(a²+b²)) (center 1.538×,
|a|=|b|=0.4 → 0.952×); heading bi lệch squirt tới ∓1.6° @|a|=0.4 (a>0 lệch
âm). ``v0_ball``/``phi_ball`` đo từ sim lưu kèm làm aux — harness baseline
cần chúng để quy đổi công bằng. ``identifiable`` = cue có ≥1 va chạm bi
hoặc băng trong sim (spin chỉ lộ qua va chạm — design §3.2).

Format ``.npz`` (đọc KHÔNG cần pooltool — BG27 train ở venv CV):
  per-shot (S,): label_v0/phi/a/b, v0_ball, phi_ball, identifiable, scratch,
    n_bb, n_cush, t_first_bb, t_first_cush (giây kể từ strike; NaN nếu
    không có), fps, upconvert, still_s, dropout_frac, n_frames, n_balls,
    shot_idx;
  ragged + offset (S+1,): xy (Σ F_i·B_i, 2 — thứ tự [frame, ball], ball 0
    LUÔN là cue), covered (Σ F_i·B_i), img_diff (Σ F_i), ball_ids (Σ B_i);
    offsets ``xy_off``/``f_off``/``b_off``.
  PTS KHÔNG lưu: t_s = arange(n_frames)/fps — CFR thật, PTS frame trùng vẫn
  hợp lệ (đúng bẫy đã đo, chỉ img_diff lộ).

Split CHỐT TRƯỚC KHI SINH: rng mỗi cú = default_rng([BASE_SEED, split, idx])
với split 0 = train (150k), 1 = held-out (5k) — hai stream độc lập, thêm
worker hay đổi shard KHÔNG đổi nội dung cú. Tái lập: chạy lại idx nào cũng
ra đúng bit (gate G-26.3 kiểm 100 cú đầu).

Chạy (venv app ``poolcoach-env`` — pooltool ở đây, CPU):
  python scripts/broadcast/gen_synth_shots.py --out "D:/Khoa luan/datasets/bb9_synth"
  # tuỳ chọn: --workers 12 --train 150000 --holdout 5000 --sanity --verify 100
"""

from __future__ import annotations

import argparse
import json
import math
import subprocess
import sys
import time
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[2]

# ------------------------------------------------------------ bảng hằng
# MỘT chỗ duy nhất — đổi số là đổi Ở ĐÂY, kèm nguồn (nếp broadcast.py).

BASE_SEED = 20260811      # ngày chốt spec — split + mọi rng con treo từ đây
SPLIT_TRAIN, SPLIT_HELDOUT = 0, 1

# Action 4D — dải đọc từ PositionPlayEnv.step (position_env.py):
#   phi = (act+1)/2·360 → [0,360); V0 = 0.5+(act+1)/2·3.5 → [0.5, 4.0];
#   a = b = act·0.4 → [−0.4, 0.4].  BRIEF nới RIÊNG V0 sample [0.5, 8] m/s
#   để phủ cả cú safety chậm (vùng cú 11) lẫn cú break.
PHI_RANGE = (0.0, 360.0)
V0_RANGE = (0.5, 8.0)
AB_MAX = 0.4

N_BALLS_RANGE = (2, 9)    # số bi mục tiêu (BRIEF: 2–9 bi + cue); luôn giữ
                          # bi 9 (nếp _gen_table eval_fullrack), subset 1..8

STILL_RANGE_S = (0.5, 1.5)   # cửa sổ đứng yên đầu clip (định nghĩa cú §5:
                             # mọi-bi-đứng-yên → đánh; HANDOFF 23 Bất ngờ 5)
SIM_MAX_S = 20.0             # cắt quỹ đạo sim (analyze cap 30s cả clip)

# --- nhiễu P0 (HANDOFF 23, số gốc; chi tiết fit tự quyết ghi tại chỗ) ---
NOISE_STILL_SIGMA_M = 0.32e-3 / math.sqrt(2 * math.log(2))
                          # jitter tĩnh: r p50 = 0.32mm (37 bi đứng yên) →
                          # σ/trục = p50/√(2ln2) ≈ 0.272mm (Gauss 2D đẳng
                          # hướng; đuôi p99 4.94mm KHÔNG mô hình — khai)
NOISE_ALONG_M_PER_MPS = 13e-3 / 2.0
                          # blur dọc hướng chạy ±13mm đo ở cú ~2 m/s (median
                          # tốc độ 10 cú P0) → fit TUYẾN TÍNH theo tốc độ
                          # (blur ∝ v·exposure), 6.5mm/(m/s)
NOISE_ALONG_MAX_M = 26e-3 # bão hoà 2× mốc đo (bbox blur không dài vô hạn —
                          # tự quyết, chưa có số đo trên 4 m/s)
NOISE_PERP_M = 0.8e-3     # vuông góc hướng chạy ~±0.8mm, không scale
V_BLUR_MIN_MPS = 0.055    # dưới ngưỡng đứng yên (V_STILL broadcast) coi như
                          # không blur — chỉ còn jitter tĩnh

FPS_CHOICES = (25, 30, 50, 60)
FPS_WEIGHTS = (0.2, 0.4, 0.2, 0.2)   # 30fps nặng hơn — phổ broadcast thật
P_UPCONVERT_30 = 0.5      # trong cú 30fps: nửa là 25fps upconvert (stream
                          # pilot chính là loại này — HANDOFF 23 Bất ngờ 1)

DROPOUT_MAX_FRAC = 0.25   # đầu cú: U(0, 0.25) thời lượng chạy của cue
GAP_RATE_PER_BALL = 0.7   # số gap ngắn ~ Poisson(0.7) mỗi bi
GAP_RANGE_S = (0.06, 0.30)

# img_diff tổng hợp — calib theo luật dedup broadcast (DUP_REL 0.3, sàn
# median 0.08) + cỡ số fixture P0 (moving ~0.5–0.8, still ~0.02):
IMG_DIFF_PER_M = 12.0     # 1 bi 2 m/s @30fps: 0.067m·12 ≈ 0.8
IMG_DIFF_STILL = 0.03     # nhiễu sensor cảnh tĩnh (< sàn 0.08 — không cờ)
IMG_DIFF_NOISE = 0.01
IMG_DIFF_DUP = 0.012      # frame trùng: gần 0 nhưng không phải 0

SHARD_SIZE = 6250         # 150k = 24 shard; held-out 5k = 1 shard

SPEC_GIT_COPY = ROOT / "scripts" / "broadcast" / "synth_spec.json"


# ---------------------------------------------- thuần numpy (unit-testable)

def frame_grid(duration_s: float, fps: int, upconvert: bool):
    """Lưới frame CFR: PTS đều ``i/fps``; nếu ``upconvert`` (25→30) thì nội
    dung lấy ở 25fps và cứ 6 frame ra có 1 frame CHÉP nội dung frame trước
    (dup) — PTS của frame trùng vẫn hợp lệ, đúng bẫy HANDOFF 23 Bất ngờ 1.

    Trả (t_out (F,), content_idx (F,) int, t_content (C,), dup_mask (F,)):
    ``content_idx[i]`` trỏ vào ``t_content`` — frame trùng trỏ CÙNG một chỉ
    số với frame trước nó."""
    n = max(int(math.floor(duration_s * fps)), 2)
    t_out = np.arange(n) / fps
    if upconvert and fps == 30:
        src = np.floor(np.arange(n) * 25.0 / 30.0).astype(np.int64)
        t_content = np.arange(src.max() + 1) / 25.0
        content_idx = src
    else:
        t_content = t_out.copy()
        content_idx = np.arange(n)
    dup_mask = np.zeros(n, dtype=bool)
    dup_mask[1:] = content_idx[1:] == content_idx[:-1]
    return t_out, content_idx, t_content, dup_mask


def ellipse_noise(rng: np.random.Generator, vel: np.ndarray) -> np.ndarray:
    """Nhiễu vị trí per-mẫu: ellipse xoay theo HƯỚNG VẬN TỐC ``vel`` (C, 2).

    σ_dọc = √(jitter² + blur_dọc(v)²), σ_ngang = √(jitter² + 0.8mm²) khi bi
    chạy (v > V_BLUR_MIN); bi đứng yên → jitter tròn 0.272mm. Mỗi mẫu độc
    lập (per-frame bbox center — đúng cách P0 đo)."""
    speed = np.hypot(vel[:, 0], vel[:, 1])
    moving = speed > V_BLUR_MIN_MPS
    blur_along = np.clip(NOISE_ALONG_M_PER_MPS * speed, 0.0,
                         NOISE_ALONG_MAX_M) * moving
    s_along = np.sqrt(NOISE_STILL_SIGMA_M ** 2 + blur_along ** 2)
    s_perp = np.where(moving,
                      math.hypot(NOISE_STILL_SIGMA_M, NOISE_PERP_M),
                      NOISE_STILL_SIGMA_M)
    e_along = rng.standard_normal(len(vel)) * s_along
    e_perp = rng.standard_normal(len(vel)) * s_perp
    with np.errstate(invalid="ignore"):
        u = np.where(moving[:, None], vel / np.maximum(speed, 1e-9)[:, None],
                     0.0)
    # bi đứng yên: hướng ellipse vô nghĩa → trục chuẩn (σ hai trục bằng nhau)
    u[~moving] = (1.0, 0.0)
    perp = np.stack([-u[:, 1], u[:, 0]], axis=1)
    return u * e_along[:, None] + perp * e_perp[:, None]


def corrupt_shot(rng: np.random.Generator, t_sim: np.ndarray,
                 xy_sim: np.ndarray, vis_until: np.ndarray, fps: int,
                 upconvert: bool, still_s: float, dropout_frac: float):
    """Quỹ đạo lý tưởng → quan sát bẩn (ba đặc trưng P0 + jitter + img_diff).

    ``t_sim`` (T,) giây kể từ STRIKE; ``xy_sim`` (T, B, 2) — bi 0 là cue;
    ``vis_until`` (B,) giây sim mà bi biến mất (vào lỗ; inf nếu còn trên
    bàn). Trả dict: t (F,), xy (F, B, 2) float32, covered (F, B) bool,
    img_diff (F,) float32, dup_mask (F,).

    Timeline clip: [0, still_s) mọi bi đứng ở vị trí đầu; strike tại
    ``still_s``; nội dung sau đó nội suy tuyến tính từ quỹ đạo sim 100Hz.
    Frame trùng CHÉP NGUYÊN vị trí đã nhiễu (cùng một lần đo) — chuỗi vận
    tốc thô vì thế dip về 0 mỗi 6 bước nếu không dedup, đúng triệu chứng
    thật."""
    T = float(t_sim[-1])
    duration = still_s + T
    t_out, content_idx, t_content, dup_mask = frame_grid(duration, fps,
                                                         upconvert)
    B = xy_sim.shape[1]
    C = len(t_content)

    # vị trí + vận tốc lý tưởng tại từng mốc NỘI DUNG (nhiễu sinh per nội
    # dung — frame trùng dùng lại nguyên realization)
    ts = np.clip(t_content - still_s, 0.0, T)
    pos = np.empty((C, B, 2))
    vel = np.zeros((C, B, 2))
    v_sim = np.gradient(xy_sim, t_sim, axis=0) if len(t_sim) > 2 else \
        np.zeros_like(xy_sim)
    pre_still = t_content < still_s
    for b_i in range(B):
        for ax in (0, 1):
            pos[:, b_i, ax] = np.interp(ts, t_sim, xy_sim[:, b_i, ax])
            vel[:, b_i, ax] = np.interp(ts, t_sim, v_sim[:, b_i, ax])
        vel[pre_still, b_i, :] = 0.0
        pos[pre_still, b_i, :] = xy_sim[0, b_i, :]

    noisy = np.empty_like(pos)
    for b_i in range(B):
        noisy[:, b_i, :] = pos[:, b_i, :] + ellipse_noise(rng, vel[:, b_i, :])

    xy_out = noisy[content_idx].astype(np.float32)          # (F, B, 2)

    # ------- visibility: bi vào lỗ biến mất; dropout đầu cú (cue); gap ngắn
    covered = np.ones((len(t_out), B), dtype=bool)
    for b_i in range(B):
        covered[:, b_i] &= (t_out - still_s) <= vis_until[b_i] + 1e-9
    # dropout đầu cú: từ strike, phủ dropout_frac × thời lượng chạy của cue
    # (shot_07: người cúi ngắm che cue quanh lúc đánh)
    if dropout_frac > 0:
        t_hide0, t_hide1 = still_s, still_s + dropout_frac * T
        covered[(t_out >= t_hide0) & (t_out < t_hide1), 0] = False
    # gap ngắn ≤0.3s rải ngẫu nhiên, độc lập từng bi
    for b_i in range(B):
        for _ in range(rng.poisson(GAP_RATE_PER_BALL)):
            g = rng.uniform(*GAP_RANGE_S)
            g0 = rng.uniform(0.0, max(duration - g, 1e-3))
            covered[(t_out >= g0) & (t_out < g0 + g), b_i] = False

    # ------- img_diff tổng hợp trên NỘI DUNG (frame trùng ~0)
    step_motion = np.zeros(len(t_out))
    d_content = np.zeros(C)
    if C > 1:
        d_content[1:] = np.abs(np.diff(pos, axis=0)).sum(axis=(1, 2))
    step_motion = d_content[content_idx]
    step_motion[dup_mask] = 0.0
    img_diff = (IMG_DIFF_PER_M * step_motion
                + np.abs(rng.normal(IMG_DIFF_STILL, IMG_DIFF_NOISE,
                                    len(t_out))))
    img_diff[dup_mask] = IMG_DIFF_DUP * rng.uniform(0.5, 1.5,
                                                    int(dup_mask.sum()))
    img_diff[0] = -1.0                       # nếp pipeline thật (frame đầu)

    return {"t": t_out.astype(np.float32), "xy": xy_out,
            "covered": covered, "img_diff": img_diff.astype(np.float32),
            "dup_mask": dup_mask}


def sample_shot_params(rng: np.random.Generator):
    """Action + fps + still + dropout cho MỘT cú — mọi lựa chọn từ ``rng``
    của cú đó (tái lập từng cú độc lập)."""
    phi = float(rng.uniform(*PHI_RANGE))
    v0 = float(rng.uniform(*V0_RANGE))
    a = float(rng.uniform(-AB_MAX, AB_MAX))
    b = float(rng.uniform(-AB_MAX, AB_MAX))
    fps = int(rng.choice(FPS_CHOICES, p=FPS_WEIGHTS))
    upconvert = bool(fps == 30 and rng.random() < P_UPCONVERT_30)
    still_s = float(rng.uniform(*STILL_RANGE_S))
    dropout = float(rng.uniform(0.0, DROPOUT_MAX_FRAC))
    return phi, v0, a, b, fps, upconvert, still_s, dropout


# ------------------------------------------------------------- tầng sim

def _sim_env():
    """Lazy import pooltool + generator thế bàn (tái dụng eval_fullrack)."""
    sys.path.insert(0, str(ROOT / "src"))
    sys.path.insert(0, str(ROOT / "scripts"))
    import pooltool as pt
    from eval_fullrack import _place_random          # generator nghiệm thu
    return pt, _place_random


def sample_board(rng, place_random, w, l, ball_r):
    """Thế bàn 2–9 bi + cue: subset 1..8 + LUÔN bi 9 (nếp _gen_table
    eval_fullrack, mở rộng đủ dải 2–9 theo BRIEF)."""
    n_obj = int(rng.integers(N_BALLS_RANGE[0], N_BALLS_RANGE[1] + 1))
    if n_obj == 9:
        ids = [str(i) for i in range(1, 10)]
    else:
        lows = sorted(rng.choice(np.arange(1, 9), size=n_obj - 1,
                                 replace=False))
        ids = [str(int(i)) for i in lows] + ["9"]
    pts = place_random(rng, w, l, len(ids) + 1, ball_r)
    return dict(zip(["cue"] + ids, pts))


def gen_one(pt, place_random, table, idx: int, split: int):
    """Sinh MỘT cú (sim + làm bẩn). Trả (meta dict, arrays dict) hoặc ném
    RuntimeError nếu 5 lần action liên tiếp đều làm pooltool chết (chưa từng
    thấy — đếm resample vào meta)."""
    rng = np.random.default_rng([BASE_SEED, split, idx])
    board = sample_board(rng, place_random, table.w, table.l,
                         float(table.balls_R) if hasattr(table, "balls_R")
                         else 0.028575)
    for attempt in range(5):
        phi, v0, a, b, fps, upconvert, still_s, dropout = \
            sample_shot_params(rng)
        system = pt.System(
            table=table,
            balls={bid: pt.Ball.create(bid, xy=tuple(xy))
                   for bid, xy in board.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)
            break
        except Exception:
            continue
    else:
        raise RuntimeError(f"shot {split}/{idx}: 5 action lien tiep loi sim")

    ball_ids = list(board.keys())            # "cue" đứng đầu (dict giữ thứ tự)
    # --- events của cue + thời điểm bi biến mất (vào lỗ)
    n_bb = n_cush = 0
    t_first_bb = t_first_cush = math.nan
    vis_until = {bid: math.inf for bid in ball_ids}
    for ev in system.events:
        et = str(ev.event_type)
        ids = tuple(getattr(ev, "ids", ()))
        if "cue" in ids:
            if et == "ball_ball":
                n_bb += 1
                if math.isnan(t_first_bb):
                    t_first_bb = float(ev.time)
            elif "cushion" in et:
                n_cush += 1
                if math.isnan(t_first_cush):
                    t_first_cush = float(ev.time)
        if "pocket" in et:
            for bid in ids:
                if bid in vis_until:
                    vis_until[bid] = min(vis_until[bid], float(ev.time))
    identifiable = (n_bb + n_cush) > 0
    scratch = math.isfinite(vis_until["cue"])
    potted_any = any(math.isfinite(vis_until[bid]) for bid in ball_ids
                     if bid != "cue")

    # --- v0/phi thật của BI (aux label — xem docstring module)
    rvw_ev, _ss, _ts = system.balls["cue"].history.vectorize()
    v_vec = rvw_ev[1, 1, :2] if len(rvw_ev) > 1 else np.zeros(2)
    v0_ball = float(np.hypot(*v_vec))
    phi_ball = float(np.degrees(np.arctan2(v_vec[1], v_vec[0])) % 360.0)

    # --- quỹ đạo continuize 100Hz, cắt SIM_MAX_S
    pt.continuize(system, dt=0.01, inplace=True)
    xy_list, t_ref = [], None
    for bid in ball_ids:
        rvw, _s2, ts2 = system.balls[bid].history_cts.vectorize()
        if t_ref is None:
            t_ref = np.asarray(ts2, dtype=np.float64)
        xy_list.append(rvw[: len(t_ref), 0, :2])
    n_keep = int(np.searchsorted(t_ref, SIM_MAX_S, side="right"))
    t_sim = t_ref[:n_keep] if n_keep >= 2 else t_ref
    xy_sim = np.stack([x[: len(t_sim)] for x in xy_list], axis=1)

    dirty = corrupt_shot(rng, t_sim, xy_sim,
                         np.array([vis_until[bid] for bid in ball_ids]),
                         fps, upconvert, still_s, dropout)

    meta = {"label_v0": v0, "label_phi": phi, "label_a": a, "label_b": b,
            "v0_ball": v0_ball, "phi_ball": phi_ball,
            "identifiable": int(identifiable), "scratch": int(scratch),
            "potted_any": int(potted_any), "n_bb": n_bb, "n_cush": n_cush,
            "t_first_bb": t_first_bb, "t_first_cush": t_first_cush,
            "fps": fps, "upconvert": int(upconvert), "still_s": still_s,
            "dropout_frac": dropout, "n_frames": len(dirty["t"]),
            "n_balls": len(ball_ids), "shot_idx": idx}
    arrays = {"xy": dirty["xy"], "covered": dirty["covered"],
              "img_diff": dirty["img_diff"],
              "ball_ids": np.array([0 if bid == "cue" else int(bid)
                                    for bid in ball_ids], dtype=np.uint8)}
    return meta, arrays


META_FIELDS = [
    ("label_v0", np.float32), ("label_phi", np.float32),
    ("label_a", np.float32), ("label_b", np.float32),
    ("v0_ball", np.float32), ("phi_ball", np.float32),
    ("identifiable", np.uint8), ("scratch", np.uint8),
    ("potted_any", np.uint8), ("n_bb", np.uint16), ("n_cush", np.uint16),
    ("t_first_bb", np.float32), ("t_first_cush", np.float32),
    ("fps", np.uint8), ("upconvert", np.uint8), ("still_s", np.float32),
    ("dropout_frac", np.float32), ("n_frames", np.uint32),
    ("n_balls", np.uint8), ("shot_idx", np.uint32),
]


def write_shard(path: Path, metas: list[dict], arrays: list[dict]) -> None:
    out = {}
    for name, dt in META_FIELDS:
        out[name] = np.array([m[name] for m in metas], dtype=dt)
    xy = np.concatenate([a["xy"].reshape(-1, 2) for a in arrays])
    covered = np.concatenate([a["covered"].reshape(-1) for a in arrays])
    img_diff = np.concatenate([a["img_diff"] for a in arrays])
    ball_ids = np.concatenate([a["ball_ids"] for a in arrays])
    nf = out["n_frames"].astype(np.int64)
    nb = out["n_balls"].astype(np.int64)
    out["xy"] = xy.astype(np.float32)
    out["covered"] = covered.astype(np.uint8)
    out["img_diff"] = img_diff.astype(np.float32)
    out["ball_ids"] = ball_ids
    out["xy_off"] = np.concatenate([[0], np.cumsum(nf * nb)]).astype(np.uint64)
    out["f_off"] = np.concatenate([[0], np.cumsum(nf)]).astype(np.uint64)
    out["b_off"] = np.concatenate([[0], np.cumsum(nb)]).astype(np.uint64)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp.npz")
    np.savez_compressed(tmp, **out)
    tmp.replace(path)


def iter_shots(shard_path: Path):
    """Đọc shard → yield dict mỗi cú (t dựng lại từ fps — PTS CFR, xem
    docstring format). KHÔNG cần pooltool — BG27/eval dùng chung loader."""
    z = np.load(shard_path)
    S = len(z["shot_idx"])
    for i in range(S):
        F = int(z["n_frames"][i])
        B = int(z["n_balls"][i])
        x0, f0, b0 = (int(z["xy_off"][i]), int(z["f_off"][i]),
                      int(z["b_off"][i]))
        shot = {name: z[name][i].item() for name, _dt in META_FIELDS}
        shot["t"] = np.arange(F) / float(z["fps"][i])
        shot["xy"] = z["xy"][x0:x0 + F * B].reshape(F, B, 2)
        shot["covered"] = z["covered"][x0:x0 + F * B].reshape(F, B) \
            .astype(bool)
        shot["img_diff"] = z["img_diff"][f0:f0 + F]
        shot["ball_ids"] = z["ball_ids"][b0:b0 + B]
        yield shot


# ------------------------------------------------------------ worker/main

def shard_path(out_dir, split: int, start: int) -> Path:
    tag = "train" if split == SPLIT_TRAIN else "heldout"
    return Path(out_dir) / f"{tag}_{start:07d}.npz"


def gen_shard(args_tuple):
    """Worker: sinh [start, start+count) của ``split`` → 1 shard riêng
    (mỗi worker ghi file mình — bẫy quen #6, không share handle)."""
    split, start, count, out_dir = args_tuple
    pt, place_random = _sim_env()
    table = pt.Table.default()
    metas, arrays = [], []
    t0 = time.time()
    for k in range(start, start + count):
        m, a = gen_one(pt, place_random, table, k, split)
        metas.append(m)
        arrays.append(a)
    path = shard_path(out_dir, split, start)
    write_shard(path, metas, arrays)
    return (path.stem, start, count, time.time() - t0,
            float(np.mean([m["identifiable"] for m in metas])))


def build_spec(n_train: int, n_holdout: int, table_w: float, table_l: float,
               ball_r: float, pooltool_version: str) -> dict:
    try:
        commit = subprocess.run(
            ["git", "-C", str(ROOT), "rev-parse", "HEAD"],
            capture_output=True, text=True, check=True).stdout.strip()
    except Exception:
        commit = "unknown"
    return {
        "name": "bb9_synth", "created": "2026-08-11", "commit": commit,
        "base_seed": BASE_SEED,
        "split_rule": "rng per shot = default_rng([base_seed, split, idx]); "
                      "split 0=train idx<n_train, 1=heldout idx<n_holdout",
        "n_train": n_train, "n_holdout": n_holdout,
        "table": {"w_m": table_w, "l_m": table_l, "ball_r_m": ball_r,
                  "source": "pooltool Table.default() — KHONG chinh physics"},
        "pooltool": pooltool_version,
        "action": {"phi_deg": list(PHI_RANGE), "v0_mps": list(V0_RANGE),
                   "ab_max": AB_MAX,
                   "v0_note": "env RL [0.5,4.0]; BRIEF noi sample [0.5,8]"},
        "board": {"n_object_balls": list(N_BALLS_RANGE),
                  "generator": "eval_fullrack._place_random + subset 1..8 "
                               "+ luon bi 9"},
        "noise": {
            "still_sigma_m": NOISE_STILL_SIGMA_M,
            "along_m_per_mps": NOISE_ALONG_M_PER_MPS,
            "along_max_m": NOISE_ALONG_MAX_M, "perp_m": NOISE_PERP_M,
            "v_blur_min_mps": V_BLUR_MIN_MPS,
            "fps_choices": list(FPS_CHOICES),
            "fps_weights": list(FPS_WEIGHTS),
            "p_upconvert_30": P_UPCONVERT_30,
            "dropout_max_frac": DROPOUT_MAX_FRAC,
            "gap_rate_per_ball": GAP_RATE_PER_BALL,
            "gap_range_s": list(GAP_RANGE_S),
            "img_diff": {"per_m": IMG_DIFF_PER_M, "still": IMG_DIFF_STILL,
                         "noise": IMG_DIFF_NOISE, "dup": IMG_DIFF_DUP},
            "source": "HANDOFF 23 Bat ngo 1-3 (P0 do that); fit chi tiet "
                      "xem docstring gen_synth_shots.py"},
        "strike_physics_probe": {
            "v_ball_over_v0": "2/(1.3 + 2.5*(a^2+b^2)) — do probe BG26b "
                              "(center 1.5385, |a|=|b|=0.4 -> 0.9524)",
            "squirt_deg_at_a04": 1.59,
            "side_sign": "a>0 -> side-L, a<0 -> side-R (probe han2005 "
                         "settle-chord, khop _read_spin docstring)"},
        "still_range_s": list(STILL_RANGE_S), "sim_max_s": SIM_MAX_S,
        "format": "npz shard; xem docstring gen_synth_shots.py (PTS = "
                  "arange(n)/fps, khong luu)",
    }


def run_sanity(out_dir: Path, split_tag: str, n_sheet: int = 20) -> dict:
    """Sanity G-26.3: histogram 4 tham số + tỷ lệ identifiable + contact
    sheet 20 cú (PNG trong ``sanity/``). Trả dict số để in/ghi HANDOFF."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    shards = sorted(out_dir.glob(f"{split_tag}_*.npz"))
    cols = {k: [] for k in ("label_v0", "label_phi", "label_a", "label_b",
                            "identifiable", "fps", "upconvert", "n_frames",
                            "scratch")}
    for p in shards:
        z = np.load(p)
        for k in cols:
            cols[k].append(z[k])
    cols = {k: np.concatenate(v) for k, v in cols.items()}
    sane = out_dir / "sanity"
    sane.mkdir(exist_ok=True)

    fig, axes = plt.subplots(2, 2, figsize=(10, 7))
    for ax, key, rng_ in zip(
            axes.flat, ["label_v0", "label_phi", "label_a", "label_b"],
            [V0_RANGE, PHI_RANGE, (-AB_MAX, AB_MAX), (-AB_MAX, AB_MAX)]):
        ax.hist(cols[key], bins=60, range=rng_)
        ax.set_title(f"{key} ({split_tag})")
    fig.tight_layout()
    fig.savefig(sane / f"hist_labels_{split_tag}.png", dpi=110)
    plt.close(fig)

    # contact sheet: 20 cú đầu shard đầu — track cue bẩn + bi khác + mép bàn
    fig, axes = plt.subplots(4, 5, figsize=(15, 13))
    for ax, shot in zip(axes.flat, iter_shots(shards[0])):
        xy, cov = shot["xy"], shot["covered"]
        for b_i in range(xy.shape[1]):
            m = cov[:, b_i]
            ax.plot(xy[m, b_i, 0], xy[m, b_i, 1],
                    ".-" if b_i == 0 else ".", ms=2, lw=0.7,
                    alpha=0.9 if b_i == 0 else 0.45)
        ax.add_patch(plt.Rectangle((0, 0), 0.9906, 1.9812, fill=False,
                                   lw=0.8))
        ax.set_title(f"#{shot['shot_idx']} v0={shot['label_v0']:.1f} "
                     f"fps={shot['fps']}{'+dup' if shot['upconvert'] else ''}"
                     f" drop={shot['dropout_frac']:.2f}", fontsize=7)
        ax.set_aspect("equal")
        ax.set_xticks([])
        ax.set_yticks([])
    fig.tight_layout()
    fig.savefig(sane / f"contact_sheet_{split_tag}.png", dpi=110)
    plt.close(fig)

    fps_mix = {int(f): int((cols["fps"] == f).sum())
               for f in np.unique(cols["fps"])}
    return {"n": int(len(cols["label_v0"])),
            "identifiable_rate": float(cols["identifiable"].mean()),
            "scratch_rate": float(cols["scratch"].mean()),
            "upconvert_rate": float(cols["upconvert"].mean()),
            "fps_mix": fps_mix,
            "v0_minmax": [float(cols["label_v0"].min()),
                          float(cols["label_v0"].max())],
            "frames_mean": float(cols["n_frames"].mean())}


def verify_repro(out_dir: Path, n: int = 100) -> bool:
    """G-26.3: sinh lại ``n`` cú đầu train từ spec (in-process) và so BIT
    với shard đã ghi."""
    pt, place_random = _sim_env()
    table = pt.Table.default()
    shard = sorted(out_dir.glob("train_*.npz"))[0]
    fresh_m, fresh_a = [], []
    for i in range(n):
        m, a = gen_one(pt, place_random, table, i, SPLIT_TRAIN)
        fresh_m.append(m)
        fresh_a.append(a)
    for i, shot in enumerate(iter_shots(shard)):
        if i >= n:
            break
        m, a = fresh_m[i], fresh_a[i]
        if not (np.array_equal(a["xy"].astype(np.float32), shot["xy"])
                and np.array_equal(a["covered"], shot["covered"])
                and np.array_equal(a["img_diff"].astype(np.float32),
                                   shot["img_diff"])
                and m["label_v0"] == np.float32(shot["label_v0"])):
            print(f"  MISMATCH tai cu {i}")
            return False
    return True


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=r"D:\Khoa luan\datasets\bb9_synth")
    ap.add_argument("--train", type=int, default=150_000)
    ap.add_argument("--holdout", type=int, default=5_000)
    ap.add_argument("--workers", type=int, default=10)
    ap.add_argument("--sanity", action="store_true")
    ap.add_argument("--verify", type=int, default=0,
                    help="so bit N cu dau train voi shard (G-26.3)")
    ap.add_argument("--only", choices=["train", "heldout", "all"],
                    default="all")
    args = ap.parse_args()
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    if args.verify:
        ok = verify_repro(out_dir, args.verify)
        print(f"verify {args.verify} cu dau: {'BIT-GIONG' if ok else 'LECH'}")
        sys.exit(0 if ok else 1)

    tasks = []
    if args.only in ("train", "all"):
        tasks += [(SPLIT_TRAIN, s, min(SHARD_SIZE, args.train - s),
                   str(out_dir))
                  for s in range(0, args.train, SHARD_SIZE)]
    if args.only in ("heldout", "all"):
        tasks += [(SPLIT_HELDOUT, s, min(SHARD_SIZE, args.holdout - s),
                   str(out_dir))
                  for s in range(0, args.holdout, SHARD_SIZE)]
    # resume idempotent: shard đã có thì bỏ qua — nội dung mỗi cú tất định
    # theo [BASE_SEED, split, idx] nên sinh lại hay giữ nguyên là một
    n_all = len(tasks)
    tasks = [t for t in tasks if not shard_path(t[3], t[0], t[1]).exists()]
    if len(tasks) < n_all:
        print(f"resume: bo qua {n_all - len(tasks)} shard da co")

    pt, _pr = _sim_env()
    import pooltool
    spec = build_spec(args.train, args.holdout,
                      float(pt.Table.default().w),
                      float(pt.Table.default().l), 0.028575,
                      getattr(pooltool, "__version__", "?"))
    (out_dir / "spec.json").write_text(json.dumps(spec, indent=2),
                                       encoding="utf-8")
    SPEC_GIT_COPY.write_text(json.dumps(spec, indent=2), encoding="utf-8")
    print(f"spec ghi: {out_dir / 'spec.json'} + {SPEC_GIT_COPY}")

    t0 = time.time()
    if args.workers <= 1:
        results = [gen_shard(t) for t in tasks]
    else:
        import multiprocessing as mp
        with mp.get_context("spawn").Pool(args.workers) as pool:
            results = []
            for r in pool.imap_unordered(gen_shard, tasks):
                results.append(r)
                stem, _start, count, dt, ident = r
                print(f"  shard {stem}: {count} cu / {dt:.0f}s "
                      f"(identifiable {ident:.1%})", flush=True)
    total = sum(r[2] for r in results)
    print(f"XONG {total} cu / {(time.time() - t0) / 60:.1f} phut")

    if args.sanity:
        for tag in ("train", "heldout"):
            if list(out_dir.glob(f"{tag}_*.npz")):
                s = run_sanity(out_dir, tag)
                print(f"sanity {tag}: {json.dumps(s)}")


if __name__ == "__main__":
    main()