#!/usr/bin/env python3
"""MEASURE LEAD CAPTURE BY RANGE: how much of the required lead do we apply?

Reads the recent LIVE recorded battles vs the real DrussGT
(``/tmp/tfil_ab2/out/<ARM>/run<N>.jsonl`` per-tick worldstate +
``.events.jsonl`` fire/hit sidecar + ``.jsonl.rounds.json`` tick index,
produced by ``tools/robocode_shim/run_bridge_battle.sh`` / ``tools/ab/ab_run.sh``)
and, for every shot WE fire, compares the lead we actually put into the bullet
against the lead the enemy's motion required -- split by RANGE band and then by
fired POWER inside each band.

WHICH SIDE IS US
----------------
``TrBattleCapture`` writes the movement subject (the bot whose name matches
``--subject``, here DrussGT) as ``e*`` and the adversary (ModularBot, our bot)
as ``s*``.  So our worldstate track is ``sx,sy,se`` and DrussGT's is
``ex,ey,ee``.  Owner ids in the event sidecar are NOT stable across runs, so
every run is attributed independently: a fire event's (x,y) is the firing
tank's centre and its energy drops by exactly the fired power one tick later.
Attribution is cross-checked against the server ``death`` events (496/496).

THE DECOMPOSITION (all angles in deg, all relative to the LOS at the FIRE tick)
------------------------------------------------------------------------------
  O        = the firing tank's centre at the fire tick (the bullet line origin)
  v        = 20 - 3p px/tick
  t*       = the AIM-INDEPENDENT interception tick: the first tick k with
             |E(t0+k) - O| <= v*k.  It depends only on the enemy's recorded
             truth and the bullet speed, never on our aim.
  required = bearing(O -> E(t*)) - LOS      (the lead a perfect gun would need)
  applied  = bullet bearing (server-recorded, perfect truth) - LOS
  leadErr  = applied - required  (wrapped)
  capture  = applied / required  (guarded: only where |required| >= 2px at range)
  miss_px  = |E(t*) - (O + v*t* * u)|       (bullet line vs enemy at t*)

CAPTURE < 1 IS NOT BY ITSELF A BUG
----------------------------------
``required`` uses the enemy's ACTUAL future dodge (perfect information), which
no real gun can know.  So the script also computes a NAIVE LINEAR PREDICTOR
control -- the same bullet speed, aimed at the intercept of a straight-line
continuation of the enemy's last LIN_M ticks of velocity.  That control is the
ceiling a trivial predictive gun reaches against this dodger; our capture is
read against it, not only against 1.0.

MEASURED = every number printed.  INFERRED = the causal reading in the doc.
"""
from __future__ import annotations

import argparse
import collections
import json
import math
import os
import statistics

US_SIDE = "s"                       # adversary = ModularBot = us
BOT_RADIUS = 18.0
SPEED_A, SPEED_B = 20.0, 3.0        # bullet speed = 20 - 3*p px/tick
MAX_FLIGHT = 220
MIN_REQ_PX = 2.0                    # capture defined only above this lateral lead
LIN_M = 4                           # ticks of velocity for the naive-predictor control

BANDS = [(0, 100), (100, 200), (200, 300), (300, 450), (450, 1e9)]
BAND_LABELS = ["0-100", "100-200", "200-300", "300-450", "450+"]
PWR_EDGES = [0.0, 0.5, 0.75, 1.0, 1.5, 100.0]
PWR_LABELS = ["<0.50", "0.50-0.75", "0.75-1.00", "1.00-1.50", ">=1.50"]


def wrap180(x):
    """Signed angular difference in (-180, 180]."""
    x = (x + 180.0) % 360.0 - 180.0
    return x + 360.0 if x <= -180.0 else x


def band_of(r):
    for (lo, hi), lab in zip(BANDS, BAND_LABELS):
        if lo <= r < hi:
            return lab
    return BAND_LABELS[-1]


def pbin(p):
    for i in range(len(PWR_EDGES) - 1):
        if PWR_EDGES[i] <= p < PWR_EDGES[i + 1]:
            return i
    return None


def mean(xs):
    return statistics.fmean(xs) if xs else float("nan")


def median(xs):
    return statistics.median(xs) if xs else float("nan")


def stdev(xs):
    return statistics.pstdev(xs) if len(xs) > 1 else float("nan")


def _rows(path):
    out = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            o = json.loads(line)
            if "tick" in o:
                out.append(o)
    return out


class Run:
    """One live battle: capture rows, event sidecar, per-round tick index."""

    def __init__(self, cap_path, events_path, rounds_path):
        self.cap_path = cap_path
        self.rows = _rows(cap_path)
        self.by_tick = {r["tick"]: r for r in self.rows}
        self.events = [json.loads(l) for l in open(events_path) if l.strip()]
        self.rounds = json.load(open(rounds_path))["rounds"]
        self.start = {r["round"]: r["startTick"] for r in self.rounds}
        self.count = {r["round"]: r["count"] for r in self.rounds}
        self.owner_side = self._resolve_owner_side()
        self.resolution = {}
        for ev in self.events:
            if ev.get("type") in ("hit", "hitwall", "hitbullet"):
                self.resolution[(ev["round"], ev["owner"], ev["bullet"])] = ev["type"]

    # -- attribution ------------------------------------------------------
    def _match(self, t, side, ev):
        r, r1 = self.by_tick.get(t), self.by_tick.get(t + 1)
        if r is None or r1 is None:
            return False
        if abs(r[side + "x"] - ev["x"]) > 0.02 or abs(r[side + "y"] - ev["y"]) > 0.02:
            return False
        ef = "ee" if side == "e" else "se"
        return abs((r[ef] - r1[ef]) - ev["power"]) < 0.02

    def _resolve_owner_side(self):
        votes = collections.defaultdict(collections.Counter)
        for ev in self.events:
            if ev.get("type") != "fire":
                continue
            guess = self.start.get(ev["round"], 0) + ev["tick"]
            for t in range(guess - 8, guess + 9):
                for side in ("e", "s"):
                    if self._match(t, side, ev):
                        votes[ev["owner"]][side] += 1
        return {o: c.most_common(1)[0][0] for o, c in votes.items() if c}

    def align(self, ev):
        guess = self.start.get(ev["round"], 0) + ev["tick"]
        side = self.owner_side.get(ev["owner"])
        best = None
        for t in range(guess - 8, guess + 9):
            for sd in (("e", "s") if side is None else (side,)):
                if self._match(t, sd, ev):
                    d = abs(t - guess)
                    if best is None or d < best[0]:
                        best = (d, t, sd)
        return None if best is None else best[1]

    def n_position_matches(self, ev):
        n = 0
        guess = self.start.get(ev["round"], 0) + ev["tick"]
        for t in range(guess - 8, guess + 9):
            for sd in ("e", "s"):
                r = self.by_tick.get(t)
                if r is None:
                    continue
                if abs(r[sd + "x"] - ev["x"]) <= 0.02 and abs(r[sd + "y"] - ev["y"]) <= 0.02:
                    n += 1
        return n

    def death_side_ok(self):
        """Every death must land on a side whose energy has collapsed to 0."""
        out = []
        for ev in self.events:
            if ev.get("type") != "death":
                continue
            side = self.owner_side.get(ev["victim"])
            g = self.start.get(ev["round"], 0) + ev["tick"]
            ekey = "ee" if side == "e" else "se"
            ok = False
            for t in range(g - 8, g + 10):
                r = self.by_tick.get(t)
                if r is not None and r[ekey] <= 1e-6:
                    ok = True
                    break
            out.append(ok)
        return out

    # -- per-shot extraction ---------------------------------------------
    def shots(self):
        for ev in self.events:
            if ev.get("type") != "fire":
                continue
            if self.owner_side.get(ev["owner"]) != US_SIDE:
                continue
            t0 = self.align(ev)
            if t0 is None:
                continue
            s = self._shot(ev, t0)
            if s is not None:
                yield s

    def _shot(self, ev, t0):
        r0 = self.by_tick[t0]
        p = ev["power"]
        v = SPEED_A - SPEED_B * p
        th = math.radians(ev["dir"])
        ux, uy = math.cos(th), math.sin(th)
        x0, y0 = ev["x"], ev["y"]                 # tank centre == bullet line origin
        rnd = ev["round"]
        end = self.start[rnd] + self.count[rnd]
        rng = math.hypot(r0["ex"] - x0, r0["ey"] - y0)

        # aim-independent interception tick (oracle)
        tstar = None
        for k in range(1, MAX_FLIGHT + 1):
            if t0 + k >= end:
                break
            r = self.by_tick.get(t0 + k)
            if r is None:
                break
            if math.hypot(r["ex"] - x0, r["ey"] - y0) <= v * k:
                tstar = (k, r)
                break
        if tstar is None:
            return None
        kstar, ea = tstar

        # miss distance: perpendicular of the enemy at the tick our (actual) bullet
        # reaches its along-track plane, relative to our bullet's real line.
        # Same definition as common_libs/tests/analyze_drussgt_dodge_vs_power.py.
        perp_arr = None
        for k in range(1, MAX_FLIGHT + 1):
            if t0 + k >= end:
                break
            r = self.by_tick.get(t0 + k)
            if r is None:
                break
            dx, dy = r["ex"] - x0, r["ey"] - y0
            al = dx * ux + dy * uy
            if al > 0 and v * k >= al:
                perp_arr = dx * uy - dy * ux
                break
        miss_px = abs(perp_arr) if perp_arr is not None else float("nan")

        los = math.degrees(math.atan2(r0["ey"] - y0, r0["ex"] - x0))
        required = wrap180(math.degrees(math.atan2(ea["ey"] - y0, ea["ex"] - x0)) - los)
        applied = wrap180(ev["dir"] - los)
        leaderr = wrap180(applied - required)

        # naive linear-predictor control: continue the last LIN_M ticks of velocity
        lin = None
        rp = self.by_tick.get(t0 - LIN_M)
        if rp is not None:
            vx = (r0["ex"] - rp["ex"]) / LIN_M
            vy = (r0["ey"] - rp["ey"]) / LIN_M
            kk = rng / v
            for _ in range(40):
                kk = math.hypot(r0["ex"] + vx * kk - x0, r0["ey"] + vy * kk - y0) / v
            lin = wrap180(math.degrees(math.atan2(r0["ey"] + vy * kk - y0,
                                                  r0["ex"] + vx * kk - x0)) - los)

        req_px = abs(math.radians(required)) * rng
        capture = (applied / required) if (req_px >= MIN_REQ_PX and abs(required) > 1e-9) else None
        capture_lin = (lin / required) if (lin is not None and req_px >= MIN_REQ_PX
                                           and abs(required) > 1e-9) else None

        kind = self.resolution.get((rnd, ev["owner"], ev["bullet"]))
        return dict(
            run=self.cap_path, rnd=rnd, tick=t0, power=p, speed=v,
            range=rng, band=band_of(rng), pbin=pbin(p), flight=kstar,
            required=required, applied=applied, leaderr=leaderr, lin=lin,
            capture=capture, capture_lin=capture_lin, req_px=req_px,
            miss_px=miss_px,
            tolerance_deg=math.degrees(math.atan(BOT_RADIUS / rng)),
            hit=1.0 if kind == "hit" else 0.0,
            resolution=kind, npos=self.n_position_matches(ev),
            our_energy=r0["se"], enemy_energy=r0["ee"], rel_tick=t0 - self.start[rnd],
        )


def discover_tfil(root):
    runs = []
    for sub in sorted(os.listdir(root)):
        d = os.path.join(root, sub)
        if not os.path.isdir(d):
            continue
        for fn in sorted(os.listdir(d)):
            if not fn.endswith(".jsonl") or fn.endswith(".events.jsonl"):
                continue
            cap = os.path.join(d, fn)
            ev, rj = cap[:-6] + ".events.jsonl", cap + ".rounds.json"
            if os.path.exists(ev) and os.path.exists(rj):
                runs.append(Run(cap, ev, rj))
    return runs


def discover_powtest(root):
    """powtest layout: cap_<arm>_r<N>.jsonl + events_<arm>_r<N>.json + .rounds.json."""
    runs = []
    if not os.path.isdir(root):
        return runs
    for fn in sorted(os.listdir(root)):
        if not (fn.startswith("cap_") and fn.endswith(".jsonl")):
            continue
        cap = os.path.join(root, fn)
        stem = fn[4:].replace(".jsonl", ".json")
        ev = os.path.join(root, "events_" + stem)
        rj = cap + ".rounds.json"
        if os.path.exists(ev) and os.path.exists(rj):
            runs.append(Run(cap, ev, rj))
    return runs


# --------------------------------------------------------------- statistics
def slope(xs, ys):
    d = sum(x * x for x in xs)
    return sum(x * y for x, y in zip(xs, ys)) / d if d > 0 else float("nan")


def cell(shots):
    caps = [s["capture"] for s in shots if s["capture"] is not None]
    lins = [s["capture_lin"] for s in shots if s["capture_lin"] is not None]
    req = [s["required"] for s in shots]
    resolved = [s for s in shots if s["resolution"] in ("hit", "hitwall", "hitbullet")]
    return dict(
        n=len(shots), nc=len(caps), excl=len(shots) - len(caps),
        unresolved=len(shots) - len(resolved),
        range=mean([s["range"] for s in shots]),
        req=mean(req),
        app=mean([s["applied"] for s in shots]),
        cap=mean(caps), cap_med=median(caps), cap_slope=slope(req, [s["applied"] for s in shots]),
        cap_lin=mean(lins), cap_lin_slope=slope(req, [s["lin"] for s in shots if s["lin"] is not None]),
        ae=mean([abs(s["leaderr"]) for s in shots]),
        ae_sd=stdev([abs(s["leaderr"]) for s in shots]),
        ae_px=mean([abs(math.radians(s["leaderr"])) * s["range"] for s in shots]),
        miss=mean([s["miss_px"] for s in shots]),
        med_miss=median([s["miss_px"] for s in shots]),
        tol=mean([s["tolerance_deg"] for s in shots]),
        hit=mean([s["hit"] for s in resolved]),
        enE=mean([s["enemy_energy"] for s in shots]),
        flight=mean([s["flight"] for s in shots]),
    )


def table(cells, title):
    print(f"\n{title}")
    hdr = (f"{'cell':<24} {'n':>6} {'excl':>5} {'range':>5} {'flt':>4} {'reqL':>7} {'appL':>7} "
           f"{'capt':>6} {'capSlp':>7} {'capLin':>7} {'|err|':>6} {'+/-':>5} {'|err|px':>7} "
           f"{'tol':>5} {'hit%':>6}")
    print(hdr)
    print("-" * len(hdr))
    for name, sh in cells:
        if not sh:
            continue
        c = cell(sh)
        print(f"{name:<24} {c['n']:>6} {c['excl']:>5} {c['range']:>5.0f} {c['flight']:>4.1f} "
              f"{c['req']:>7.2f} {c['app']:>7.2f} {c['cap']:>6.2f} {c['cap_slope']:>7.3f} "
              f"{c['cap_lin_slope']:>7.3f} {c['ae']:>6.2f} {c['ae_sd']:>5.2f} {c['ae_px']:>7.1f} "
              f"{c['tol']:>5.2f} {100*c['hit']:>5.2f}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tfil", default="/tmp/tfil_ab2/out")
    ap.add_argument("--powtest", default="/tmp/powtest")
    ap.add_argument("--json", default=None)
    args = ap.parse_args()

    runs = discover_tfil(args.tfil)
    print("=" * 132)
    print("LEAD CAPTURE BY RANGE  --  live ModularBot vs real DrussGT")
    print("=" * 132)
    print(f"corpus            : {args.tfil}")
    print(f"battles (runs)    : {len(runs)}")
    if not runs:
        raise SystemExit("no runs found")

    amb = 0
    death_ok = death_n = 0
    pow_all = collections.defaultdict(collections.Counter)
    sideowners = collections.Counter()
    allshots = []
    for rn in runs:
        for o, sd in rn.owner_side.items():
            sideowners[(sd, o)] += 1
        amb += sum(1 for ev in rn.events if ev.get("type") == "fire"
                   and rn.n_position_matches(ev) > 1)
        for ev in rn.events:
            if ev.get("type") == "fire":
                pow_all[rn.owner_side.get(ev["owner"], "?")][round(ev["power"], 2)] += 1
        ds = rn.death_side_ok()
        death_n += len(ds)
        death_ok += sum(ds)
        allshots.extend(rn.shots())

    sides = collections.Counter(tuple(sorted(rn.owner_side.values())) for rn in runs)
    print(f"per-run owner maps: {dict(sides)}   (expect all ('e','s'))")
    print(f"runs with a broken owner map      : {sum(1 for rn in runs if sorted(rn.owner_side.values()) != ['e','s'])}")
    print(f"fire events w/ ambiguous position : {amb}")
    print(f"death events side-validated       : {death_ok}/{death_n}")
    for sd in ("s", "e"):
        tot = sum(pow_all[sd].values())
        top = pow_all[sd].most_common(6)
        print(f"  side {sd} = {'US (ModularBot)' if sd == US_SIDE else 'DrussGT'}  "
              f"n={tot:>6}  top powers: {top}")

    hits = [s for s in allshots if s["hit"] == 1.0]
    misses = [s for s in allshots if s["resolution"] in ("hitwall", "hitbullet")
              and s["hit"] == 0.0]
    unresolved = [s for s in allshots if s["resolution"] is None]
    print("\n" + "=" * 132)
    print("SAMPLE")
    print("=" * 132)
    print(f"our shots (aligned, with an interception)     : {len(allshots)}")
    print(f"  server hit / hitwall / hitbullet / other    : "
          f"{sum(1 for s in allshots if s['resolution']=='hit')} / "
          f"{sum(1 for s in allshots if s['resolution']=='hitwall')} / "
          f"{sum(1 for s in allshots if s['resolution']=='hitbullet')} / "
          f"{len(unresolved)}")
    print(f"capture excluded (|required| < {MIN_REQ_PX} px lateral) : "
          f"{sum(1 for s in allshots if s['capture'] is None)}")

    print("\n" + "=" * 132)
    print("GEOMETRY VALIDATION  (hits must separate from misses)")
    print("=" * 132)
    hdr = f"{'':<16}{'n':>7} {'|leadErr|deg':>13} {'|leadErr|px':>12} {'misPx mean':>11} {'misPx med':>10} {'<18px':>7}"
    print(hdr)
    for lab, grp in (("HITS", hits), ("MISSES", misses)):
        if not grp:
            continue
        ae = [abs(s["leaderr"]) for s in grp]
        apx = [abs(math.radians(s["leaderr"])) * s["range"] for s in grp]
        mp = [s["miss_px"] for s in grp]
        print(f"{lab:<16}{len(grp):>7} {mean(ae):>13.3f} {mean(apx):>12.1f} "
              f"{mean(mp):>11.1f} {median(mp):>10.1f} {100*sum(1 for x in mp if x<18)/len(mp):>6.1f}%")
    if hits and misses:
        sep = mean([s["miss_px"] for s in misses]) / max(1e-9, mean([s["miss_px"] for s in hits]))
        print(f"\nmiss separation ratio (misses/hits) = {sep:.2f}x  "
              f"-> {'OK (bearing recovery validated)' if sep > 2 else 'WEAK - bearing recovery suspect'}")
    print(f"overall measured hit rate = {100*mean([s['hit'] for s in allshots]):.2f}%")

    cells = [(lab, [s for s in allshots if s["band"] == lab]) for lab in BAND_LABELS]
    print("\n" + "=" * 132)
    print("MAIN TABLE  --  all our shots, by RANGE band")
    print("=" * 132)
    print("reqL/appL = mean signed lead vs LOS (deg).  capt = mean ratio, capSlp = robust proportional")
    print("capture, capLin = the same slope for a naive linear predictor (control).")
    table(cells, "by range band")

    print("\n" + "=" * 132)
    print("POWER WITHIN RANGE BAND  --  this is what separates 'range' from 'power'")
    print("=" * 132)
    for lab in BAND_LABELS:
        sub = [s for s in allshots if s["band"] == lab]
        if not sub:
            continue
        pcs = [(f"{lab} {PWR_LABELS[i]}", [s for s in sub if s["pbin"] == i])
               for i in range(len(PWR_LABELS))]
        table(pcs, f"range {lab}")

    # ---- endgame confound and its removal --------------------------------
    print("\n" + "=" * 132)
    print("CONFOUND CHECK  --  the sub-0.5-power shots at long range are the ENDGAME finish")
    print("=" * 132)
    print("hit% of our 450+ shots, by the ENEMY's remaining energy at the fire tick:")
    sub = [s for s in allshots if s["band"] == "450+"]
    for lo, hi in ((0, 2), (2, 5), (5, 10), (10, 20), (20, 40), (40, 150)):
        g = [s for s in sub if lo <= s["enemy_energy"] < hi]
        if g:
            c = cell(g)
            print(f"  enemy energy [{lo:>3},{hi:>3})  n={c['n']:>6}  meanP={mean([x['power'] for x in g]):>5.2f} "
                  f"capSlp={c['cap_slope']:>6.3f}  hit%={100*c['hit']:>5.2f}")
    print("\nsame table, EXCLUDING shots fired while the enemy still has >= 5 energy")
    print("(this removes the near-dead-target endgame and leaves the clean power comparison):")
    for lab in BAND_LABELS:
        sub = [s for s in allshots if s["band"] == lab and s["enemy_energy"] >= 5.0]
        if not sub:
            continue
        pcs = [(f"{lab} {PWR_LABELS[i]}", [s for s in sub if s["pbin"] == i])
               for i in range(len(PWR_LABELS))]
        table(pcs, f"range {lab}  (enemy energy >= 5)")

    # ---- response curve ---------------------------------------------------
    print("\n" + "=" * 132)
    print("RESPONSE CURVE  --  do we move the aim in proportion to the required lead?")
    print("=" * 132)
    for lab in ("200-300", "300-450", "450+"):
        g = sorted([s for s in allshots if s["band"] == lab], key=lambda s: s["required"])
        if len(g) < 100:
            continue
        nb = 8
        print(f"\nrange {lab}  (deciles of required lead -> mean applied lead):")
        for i in range(nb):
            q = g[i * len(g) // nb:(i + 1) * len(g) // nb]
            c = cell(q)
            print(f"  required [{q[0]['required']:>7.2f},{q[-1]['required']:>7.2f}]  "
                  f"meanReq={c['req']:>7.2f}  meanApp={c['app']:>7.2f}  "
                  f"hit%={100*c['hit']:>5.2f}  n={c['n']:>5}")

    print("\n" + "=" * 132)
    print("DIRECT ANSWER  --  does capture degrade with range / power?")
    print("=" * 132)
    print(f"{'range band':<12} {'n':>6} {'capSlp':>7} {'capLin':>7} {'us/lin':>7} "
          f"{'medcap':>7} {'|req|':>6} {'|app|':>6} {'|err|px':>8} {'tol':>5} {'|err|/tol':>9} {'hit%':>6}")
    for lab, sh in cells:
        if not sh:
            continue
        c = cell(sh)
        ratio = c["ae_px"] / (c["tol"] * math.pi / 180.0 * c["range"])
        areq = mean([abs(s["required"]) for s in sh])
        aapp = mean([abs(s["applied"]) for s in sh])
        print(f"{lab:<12} {c['n']:>6} {c['cap_slope']:>7.3f} {c['cap_lin_slope']:>7.3f} "
              f"{c['cap_slope']/c['cap_lin_slope'] if c['cap_lin_slope'] else float('nan'):>7.2f} "
              f"{c['cap_med']:>7.3f} {areq:>6.2f} {aapp:>6.2f} {c['ae_px']:>8.1f} {c['tol']:>5.2f} "
              f"{ratio:>9.2f} {100*c['hit']:>5.2f}")
    print("\n(capSlp = our proportional lead capture; capLin = naive linear predictor control;")
    print(" us/lin = our capture / the trivial-predictor ceiling; |req|/|app| = mean ABSOLUTE")
    print(" lead magnitude in deg; |err|/tol = mean lead error px / 18px)")
    print("\nPOWER WITHIN 450+ BAND, enemy energy >= 5 (removes the endgame confound):")
    sub = [s for s in allshots if s["band"] == "450+" and s["enemy_energy"] >= 5.0]
    for i in range(len(PWR_LABELS)):
        g = [s for s in sub if s["pbin"] == i]
        if not g:
            continue
        c = cell(g)
        print(f"  power {PWR_LABELS[i]:<11} n={c['n']:>6} meanP={mean([s['power'] for s in g]):>5.2f} "
              f"meanRange={c['range']:>4.0f} capSlp={c['cap_slope']:>6.3f} "
              f"capLin={c['cap_lin_slope']:>6.3f} |err|px={c['ae_px']:>6.1f} hit%={100*c['hit']:>5.2f}")

    if args.powtest:
        pruns = discover_powtest(args.powtest)
        if pruns:
            pshots = [s for rn in pruns for s in rn.shots()]
            print("\n" + "=" * 132)
            print(f"ROBUSTNESS CROSS-CHECK  --  powtest corpus ({len(pruns)} runs, "
                  f"{len(pshots)} shots, DIFFERENT binary)")
            print("=" * 132)
            pc = [(lab, [s for s in pshots if s["band"] == lab]) for lab in BAND_LABELS]
            table(pc, "by range band (powtest)")

    if args.json:
        blob = dict(
            corpus=args.tfil, runs=len(runs), shots=len(allshots),
            hits=len(hits), misses=len(misses),
            excl=sum(1 for s in allshots if s["capture"] is None),
            death_ok=death_ok, death_n=death_n, amb=amb,
            overall_hit=mean([s["hit"] for s in allshots]),
            hit_ae=mean([abs(s["leaderr"]) for s in hits]),
            miss_ae=mean([abs(s["leaderr"]) for s in misses]),
            hit_miss_px=mean([s["miss_px"] for s in hits]),
            miss_miss_px=mean([s["miss_px"] for s in misses]),
            pwr_hist={sd: dict(pow_all[sd]) for sd in pow_all},
            bands={lab: cell([s for s in allshots if s["band"] == lab]) for lab in BAND_LABELS},
            bands_enemy_alive={lab: cell([s for s in allshots if s["band"] == lab
                                         and s["enemy_energy"] >= 5.0])
                               for lab in BAND_LABELS},
            band_power_alive={lab: {PWR_LABELS[i]: cell([s for s in allshots
                                                         if s["band"] == lab
                                                         and s["pbin"] == i
                                                         and s["enemy_energy"] >= 5.0])
                                    for i in range(len(PWR_LABELS))}
                              for lab in BAND_LABELS},
            hit_by_enemy_energy={f"{lo}-{hi}": cell([s for s in allshots
                                                     if s["band"] == "450+"
                                                     and lo <= s["enemy_energy"] < hi])
                                 for lo, hi in ((0, 2), (2, 5), (5, 10), (10, 20), (20, 40), (40, 150))},
            band_power={lab: {PWR_LABELS[i]: cell([s for s in allshots
                                                   if s["band"] == lab and s["pbin"] == i])
                              for i in range(len(PWR_LABELS))} for lab in BAND_LABELS},
        )
        with open(args.json, "w") as f:
            json.dump(blob, f, indent=2)
        print(f"\n[json] wrote {args.json}")


if __name__ == "__main__":
    main()