#!/usr/bin/env python3
"""ab_range_bands.py — hit rate BY RANGE BAND per arm, from an ab_run.sh session.

    python3 tools/ab/ab_range_bands.py <session_dir> [--reference ARM]

Why this exists: ab_analyze.py judges damage/run and round wins over the whole
battle, but the HeadOn-vs-Pattern claim is specifically about LONG RANGE. A
whole-battle win that comes only from close range would NOT support it, so the
range breakdown is the load-bearing view. ab_analyze.py does not emit it.

Method: reuse `Run` from common_libs/tests/analyze_lead_capture_by_range.py,
which aligns each of OUR fire events to the capture tick (owner side resolved
per run, round-startTick + local-tick search) and gives the fire-tick range
(shooter->target distance) plus the server-resolved outcome of that bullet.
A band's hit rate is then hits/shots over that band's shots, pooled across the
arm's runs. A two-sided permutation test on the PER-RUN band rates (only runs
with shots in the band) is printed alongside, so a pooled difference can be
told apart from run-to-run noise; the MDE for those n is printed too.

Bands match the offline ruler (prediction_quality.nim) and the lead-capture
report: 0-100, 100-200, 200-300, 300-450, 450+ px.
"""
import argparse
import collections
import math
import os
import random
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(_HERE, "..", "..", "common_libs", "tests"))
from analyze_lead_capture_by_range import Run, BAND_LABELS  # noqa: E402

Z_ALPHA_POWER = 1.959963984540054 + 0.8416212335729143
MC_SEED = 0x5EED5EED
MC_DRAWS = 200_000


def load_arm(session_dir, arm):
    armdir = os.path.join(session_dir, arm)
    runs = []
    for fn in sorted(os.listdir(armdir)):
        if not fn.endswith(".jsonl") or fn.endswith(".events.jsonl"):
            continue
        cap = os.path.join(armdir, fn)
        ev, rj = cap[:-6] + ".events.jsonl", cap + ".rounds.json"
        if os.path.exists(ev) and os.path.exists(rj):
            runs.append((fn, cap, ev, rj))
    out = {}
    for fn, cap, ev, rj in runs:
        try:
            r = Run(cap, ev, rj)
        except Exception as e:  # noqa: BLE001
            print(f"  WARN {arm}/{fn}: {e}", file=sys.stderr)
            continue
        shots = list(r.shots())
        out[fn] = shots
    return out


def perm_test(xa, xb):
    """Two-sided permutation on the difference of means (exact if small)."""
    na, nb = len(xa), len(xb)
    if na == 0 or nb == 0:
        return None
    import itertools
    obs = abs(sum(xa) / na - sum(xb) / nb)
    pooled = list(xa) + list(xb)
    n = na + nb
    total = sum(pooled)
    ncomb = math.comb(n, na)
    if ncomb <= 20_000_000:
        cnt = 0
        for combo in itertools.combinations(range(n), na):
            sa = sum(pooled[i] for i in combo)
            if abs(sa / na - (total - sa) / nb) >= obs - 1e-9:
                cnt += 1
        return obs, cnt / ncomb, "exact", 0.0
    rng = random.Random(MC_SEED)
    B = MC_DRAWS
    cnt = 0
    for _ in range(B):
        sa = 0
        for i in rng.sample(range(n), na):
            sa += pooled[i]
        if abs(sa / na - (total - sa) / nb) >= obs - 1e-9:
            cnt += 1
    p = (cnt + 1) / (B + 1)
    return obs, p, f"MC/B={B:,}", math.sqrt(p * (1 - p) / (B + 1))


def _sd(xs):
    m = sum(xs) / len(xs)
    return math.sqrt(sum((v - m) ** 2 for v in xs) / (len(xs) - 1))


def banded(shots):
    """{band: (shots, hits)}, plus overall."""
    c = collections.Counter()
    h = collections.Counter()
    for s in shots:
        c[s["band"]] += 1
        h[s["band"]] += int(s["hit"])
    c["ALL"] += len(shots)
    h["ALL"] += int(sum(s["hit"] for s in shots))
    return c, h


def per_run_band_rate(runs, band):
    """Per-run hit rate for a band, only runs that fired in that band."""
    out = []
    for fn, shots in runs.items():
        sub = shots if band == "ALL" else [s for s in shots if s["band"] == band]
        if sub:
            out.append(sum(s["hit"] for s in sub) / len(sub))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("session_dir")
    ap.add_argument("--reference", default=None)
    args = ap.parse_args()

    arms = [d for d in sorted(os.listdir(args.session_dir))
            if os.path.isdir(os.path.join(args.session_dir, d))
            and not d.startswith(".") and d != "frozen"]
    arms = [a for a in arms if any(
        f.endswith(".jsonl") and not f.endswith(".events.jsonl")
        for f in os.listdir(os.path.join(args.session_dir, a)))]
    ref = args.reference or arms[0]
    data = {a: load_arm(args.session_dir, a) for a in arms}

    print("=" * 100)
    print("HIT RATE BY RANGE BAND  (our shots; band = shooter->target distance px at the fire tick)")
    print(f"session: {args.session_dir}   arms: {', '.join(arms)}   reference: {ref}")
    print("=" * 100)
    labels = BAND_LABELS + ["ALL"]
    hdr = f"{'band':<10}" + "".join(f"{a:>26}" for a in arms)
    print(hdr)
    print("-" * len(hdr))
    for band in labels:
        cells = []
        for a in arms:
            c, h = banded([s for shots in data[a].values() for s in shots])
            n = c[band]
            cells.append(f"{n:>7} {h[band]:>5} {(100*h[band]/n if n else 0):>6.1f}%")
        print(f"{band:<10}" + "".join(f"{c:>26}" for c in cells))

    print("\nPER-RUN BAND RATES (shows the spread behind the pooled numbers)")
    for band in labels:
        print(f"  band {band}")
        for a in arms:
            rates = per_run_band_rate(data[a], band)
            txt = " ".join(f"{100*r:.0f}" for r in rates)
            print(f"    {a:<10} n={len(rates):>2}  {txt}")

    print("\nPER-BAND PERMUTATION TEST vs `" + ref + "` (per-run rates, two-sided)")
    print(f"{'band':<10} {'arm':<10} {'d(pp)':>7} {'p':>9} {'method':<12} {'MCse':>7} {'MDE(pp)':>8}")
    print("-" * 66)
    for band in labels:
        xa = per_run_band_rate(data[ref], band)
        if len(xa) < 2:
            continue
        sd = _sd(xa)
        mde = Z_ALPHA_POWER * sd * math.sqrt(2.0 / len(xa)) * 100
        for a in arms:
            if a == ref:
                continue
            xb = per_run_band_rate(data[a], band)
            res = perm_test(xa, xb)
            if res is None:
                continue
            obs, p, method, se = res
            signed = (sum(xb) / len(xb) - sum(xa) / len(xa)) * 100  # arm - ref
            print(f"{band:<10} {a:<10} {signed:>+7.2f} {p:>9.4f} {method:<12} "
                  f"{se:>7.4f} {mde:>8.2f}")
    return 0


if __name__ == "__main__":
    sys.exit(main())