#!/usr/bin/env python3
"""ab_mechanism.py — MECHANISM CHECK for a movement A/B session (ab_run.sh).

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

ab_analyze.py answers "did the arm win more?", this answers "did the arm
actually MOVE differently?" — because a win/loss difference is uninterpretable
if the knob never changed the trajectory.

Three measurements, all from the captured per-tick worldstate (both bots'
positions every tick) plus the fire/hit event sidecar:

1. CENTRAL-BOX OCCUPANCY. Fraction of ticks our tank spends inside the
   144x144 px box centred on the arena centre (x 328-472, y 228-372) — the box
   the removed virtual pillar covered. Spawns are bottom-left (us) / top (enemy),
   never in the box, so any occupancy is genuine transit.

2. DISTANCE TO THE ENEMY. Per-tick shooter->target distance: mean, p10/median/
   p90 and the fraction of ticks in 0-100 / 100-200 / 200-300 / 300-400 / 400+ px.
   The flat heat model was measured to keep us at ~400 px; a change that works
   should move this distribution.

3. BULLET PROXIMITY. Reconstruct every bullet from its fire event (x, y, dir,
   power -> speed = 20-3p) and the resolving hit/hitwall/hitbullet event, and
   take the per-tick distance from our tank to the NEAREST live bullet. Reported
   as mean min-distance and the fraction of ticks with a bullet within 50/100/150
   px. This is how much time we spend near live bullet paths.

All three are aggregated PER ROUND and then averaged over rounds (rounds differ
in length and in how early somebody dies, so a pooled mean would weight a long
lost round more than a short won one), and a two-sided permutation test on the
analysed runs' per-run means says whether an arm's behaviour differs from the
reference beyond run-to-run noise.

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

import argparse
import itertools
import json
import math
import os
import random
import sys

ARENA_W, ARENA_H = 800.0, 600.0
BOX_X0, BOX_X1 = 328.0, 472.0     # 144x144 px centred on (400, 300)
BOX_Y0, BOX_Y1 = 228.0, 372.0
DIST_BANDS = [(0, 100), (100, 200), (200, 300), (300, 400), (400, 1e9)]
DIST_LABELS = ["0-100", "100-200", "200-300", "300-400", "400+"]
NEAR_PX = [50.0, 100.0, 150.0]
SPEED_A, SPEED_B = 20.0, 3.0
US_SIDE = "s"                     # adversary (ModularBot) = s* (see the capture)
MC_SEED = 0x5EED5EED
MC_DRAWS = 200_000
Z_ALPHA_POWER = 1.959963984540054 + 0.8416212335729143


def rows_of(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


def events_of(path):
    return [json.loads(l) for l in open(path) if l.strip()]


def resolve_owner_side(rows_by_tick, rounds, events):
    """owner id -> 's' (us) / 'e' (DrussGT), by matching each fire event's
    (x, y) to a capture row within +-8 ticks of `startTick + ev.tick`."""
    start = {r["round"]: r["startTick"] for r in rounds}
    votes = {}
    for ev in events:
        if ev.get("type") != "fire":
            continue
        g = start.get(ev["round"], 0) + ev["tick"]
        for t in range(g - 8, g + 9):
            r = rows_by_tick.get(t)
            if r is None:
                continue
            for side in ("s", "e"):
                if (abs(r[side + "x"] - ev["x"]) <= 0.02
                        and abs(r[side + "y"] - ev["y"]) <= 0.02):
                    d = votes.setdefault(ev["owner"], {"s": 0, "e": 0})
                    d[side] += 1
    return {o: ("s" if d["s"] >= d["e"] else "e") for o, d in votes.items()}


def analyse_run(cap_path, ev_path, rj_path):
    rows = rows_of(cap_path)
    if not rows:
        return None
    by_tick = {r["tick"]: r for r in rows}
    rounds = json.load(open(rj_path))["rounds"]
    events = events_of(ev_path)
    oside = resolve_owner_side(by_tick, rounds, events)

    # bullet tracks: (round, owner, bullet) -> start tick / direction / speed.
    # A live window is [t0, resolving event] (hit / hitwall / hitbullet); a
    # bullet with no resolver (round ended mid-flight) is bounded at t0 + 400.
    start = {r["round"]: r["startTick"] for r in rounds}
    bullets = {}
    for ev in events:
        key = (ev["round"], ev.get("owner"), ev.get("bullet"))
        if ev.get("type") == "fire":
            p = ev["power"]
            v = SPEED_A - SPEED_B * p
            th = math.radians(ev["dir"])
            bullets[key] = {
                "t0": start[ev["round"]] + ev["tick"],
                "x": ev["x"], "y": ev["y"],
                "ux": math.cos(th), "uy": math.sin(th), "v": v,
                "side": oside.get(ev["owner"]), "end": None,
            }
        elif ev.get("type") in ("hit", "hitwall", "hitbullet"):
            if key in bullets:
                bullets[key]["end"] = start[ev["round"]] + ev["tick"]

    # per-round accumulators
    per_round = []
    for rd in rounds:
        a, n = rd["startTick"], rd["count"]
        ticks = [t for t in range(a, a + n) if t in by_tick]
        if not ticks:
            continue
        box = near = near_e = 0
        dsum = 0.0
        mindist_sum = 0.0
        mindist_e_sum = 0.0
        band_cnt = [0] * len(DIST_BANDS)
        near_cnt = [0] * len(NEAR_PX)
        near_e_cnt = [0] * len(NEAR_PX)
        live = [(k, b) for k, b in bullets.items() if k[0] == rd["round"]]
        for t in ticks:
            r = by_tick[t]
            sx, sy = r[US_SIDE + "x"], r[US_SIDE + "y"]
            ox, oy = r["ex"], r["ey"]
            if BOX_X0 <= sx <= BOX_X1 and BOX_Y0 <= sy <= BOX_Y1:
                box += 1
            d = math.hypot(ox - sx, oy - sy)
            dsum += d
            for i, (lo, hi) in enumerate(DIST_BANDS):
                if lo <= d < hi:
                    band_cnt[i] += 1
                    break
            md = None
            mde_ = None
            for _k, b in live:
                if b["t0"] is None or t < b["t0"]:
                    continue
                te = b["end"] if b["end"] is not None else b["t0"] + 400
                if t > te:
                    continue
                dt = t - b["t0"]
                bx = b["x"] + b["v"] * dt * b["ux"]
                by = b["y"] + b["v"] * dt * b["uy"]
                if not (0.0 <= bx <= ARENA_W and 0.0 <= by <= ARENA_H):
                    continue
                dd = math.hypot(bx - sx, by - sy)
                if md is None or dd < md:
                    md = dd
                if b["side"] == "e" and (mde_ is None or dd < mde_):
                    mde_ = dd
            if md is not None:
                near += 1
                mindist_sum += md
                for i, x in enumerate(NEAR_PX):
                    if md <= x:
                        near_cnt[i] += 1
            if mde_ is not None:
                near_e += 1
                mindist_e_sum += mde_
                for i, x in enumerate(NEAR_PX):
                    if mde_ <= x:
                        near_e_cnt[i] += 1
        nt = len(ticks)
        per_round.append({
            "round": rd["round"], "ticks": nt,
            "box_frac": box / nt,
            "mean_dist": dsum / nt,
            "dist_bands": [c / nt for c in band_cnt],
            "mean_minbullet": (mindist_sum / near if near else float("nan")),
            "near_frac": [c / nt for c in near_cnt],
            "mean_minbullet_enemy": (mindist_e_sum / near_e if near_e
                                     else float("nan")),
            "near_enemy_frac": [c / nt for c in near_e_cnt],
            "ticks_with_bullet_frac": near / nt,
            "ticks_with_enemy_bullet_frac": near_e / nt,
        })
    return {"cap": cap_path, "per_round": per_round, "owner_side": oside,
            "n_rounds": len(per_round)}


def _mean(xs):
    return sum(xs) / len(xs) if xs else float("nan")


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


def perm_test(xa, xb):
    na, nb = len(xa), len(xb)
    if na == 0 or nb == 0:
        return None
    obs = abs(_mean(xa) - _mean(xb))
    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"
    rng = random.Random(MC_SEED)
    cnt = 0
    for _ in range(MC_DRAWS):
        sa = sum(pooled[i] for i in rng.sample(range(n), na))
        if abs(sa / na - (total - sa) / nb) >= obs - 1e-9:
            cnt += 1
    return obs, (cnt + 1) / (MC_DRAWS + 1), f"MC/B={MC_DRAWS:,}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("session_dir")
    ap.add_argument("--reference", default=None)
    args = ap.parse_args()
    root = args.session_dir
    arms = [d for d in sorted(os.listdir(root))
            if os.path.isdir(os.path.join(root, d)) and d != "frozen"
            and not d.startswith(".") and any(
                f.endswith(".jsonl") and not f.endswith(".events.jsonl")
                for f in os.listdir(os.path.join(root, d)))]
    ref = args.reference or arms[0]
    data = {}
    for arm in arms:
        d = os.path.join(root, arm)
        runs = []
        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):
                a = analyse_run(cap, ev, rj)
                if a:
                    runs.append(a)
        data[arm] = runs

    print("=" * 100)
    print("MECHANISM CHECK — did the arm actually move differently?")
    print(f"session {root}  arms {', '.join(arms)}  reference {ref}")
    print("box = 144x144 px centred on the arena centre (x 328-472, y 228-372)")
    print("=" * 100)

    metrics = [
        ("box_frac", "central-box occupancy", 100.0, "%"),
        ("mean_dist", "mean dist to enemy", 1.0, "px"),
        ("near_enemy_frac[1]", "ticks enemy bullet<=100px", 100.0, "%"),
        ("mean_minbullet_enemy", "mean dist nearest enemy bullet", 1.0, "px"),
        ("near_frac[1]", "ticks ANY bullet<=100px", 100.0, "%"),
    ]

    def per_run(arm, key, band=None):
        out = []
        for a in data[arm]:
            vals = []
            for pr in a["per_round"]:
                if key.startswith("near_enemy_frac["):
                    vals.append(pr["near_enemy_frac"][int(key[16])])
                elif key.startswith("near_frac["):
                    vals.append(pr["near_frac"][int(key[10])])
                else:
                    vals.append(pr[key])
            out.append(_mean(vals))
        return out

    print("\nPER-ARM MECHANISM (mean of the per-round values, then over runs)")
    hdr = (f"{'arm':<10} {'runs':>4} {'box%':>7} {'dist px':>8} "
           f"{'p10':>6} {'med':>6} {'p90':>6} {'enB<=100px%':>13} "
           f"{'enBmin px':>11}")
    print(hdr)
    print("-" * len(hdr))
    for arm in arms:
        runs = data[arm]
        if not runs:
            print(f"{arm:<10} {0:>4}  (no runs)")
            continue
        boxf = [_mean([pr["box_frac"] for pr in a["per_round"]]) for a in runs]
        md = [_mean([pr["mean_dist"] for pr in a["per_round"]]) for a in runs]
        nb = [_mean([pr["near_enemy_frac"][1] for pr in a["per_round"]])
              for a in runs]
        mb = [_mean([pr["mean_minbullet_enemy"] for pr in a["per_round"]])
              for a in runs]
        allmd = sorted(v for a in runs for v in
                       [pr["mean_dist"] for pr in a["per_round"]])
        p10 = allmd[int(0.10 * (len(allmd) - 1))]
        med = allmd[len(allmd) // 2]
        p90 = allmd[int(0.90 * (len(allmd) - 1))]
        print(f"{arm:<10} {len(runs):>4} {100*_mean(boxf):>7.2f} "
              f"{_mean(md):>8.1f} {p10:>6.0f} {med:>6.0f} {p90:>6.0f} "
              f"{100*_mean(nb):>13.2f} {_mean(mb):>11.1f}")

    print("\nDISTANCE-TO-ENEMY BANDS (fraction of ticks, mean over rounds/runs, %)")
    hdr = f"{'arm':<10}" + "".join(f"{lb:>10}" for lb in DIST_LABELS)
    print(hdr)
    print("-" * len(hdr))
    for arm in arms:
        runs = data[arm]
        if not runs:
            continue
        cells = []
        for i in range(len(DIST_BANDS)):
            v = _mean([_mean([pr["dist_bands"][i] for pr in a["per_round"]])
                       for a in runs])
            cells.append(f"{100*v:>10.2f}")
        print(f"{arm:<10}" + "".join(cells))

    print("\nENEMY-BULLET PROXIMITY (nearest LIVE ENEMY bullet; our own bullets are "
          "excluded because a bullet is born at its own tank, so 'any bullet' "
          "is dominated by our just-fired shot)")
    hdr = (f"{'arm':<10}" + "".join(f"{('<=%gpx' % x):>9}" for x in NEAR_PX)
           + f"{'min px':>9}{'enemy%':>9}")
    print(hdr)
    print("-" * len(hdr))
    for arm in arms:
        runs = data[arm]
        if not runs:
            continue
        cells = []
        for i in range(len(NEAR_PX)):
            v = _mean([_mean([pr["near_enemy_frac"][i] for pr in a["per_round"]])
                       for a in runs])
            cells.append(f"{100*v:>9.2f}")
        mb = _mean([_mean([pr["mean_minbullet_enemy"] for pr in a["per_round"]])
                    for a in runs])
        anyv = _mean([_mean([pr["ticks_with_enemy_bullet_frac"]
                             for pr in a["per_round"]]) for a in runs])
        print(f"{arm:<10}" + "".join(cells) + f"{mb:>9.1f}{100*anyv:>9.2f}")

    print("\nANY-BULLET PROXIMITY (per the task text: nearest live bullet, ours "
          "included; context only)")
    hdr = (f"{'arm':<10}" + "".join(f"{('<=%gpx' % x):>9}" for x in NEAR_PX)
           + f"{'min px':>9}{'any%':>9}")
    print(hdr)
    print("-" * len(hdr))
    for arm in arms:
        runs = data[arm]
        if not runs:
            continue
        cells = []
        for i in range(len(NEAR_PX)):
            v = _mean([_mean([pr["near_frac"][i] for pr in a["per_round"]])
                       for a in runs])
            cells.append(f"{100*v:>9.2f}")
        anyv = _mean([_mean([pr["ticks_with_bullet_frac"] for pr in a["per_round"]])
                      for a in runs])
        mb = _mean([_mean([pr["mean_minbullet"] for pr in a["per_round"]])
                    for a in runs])
        print(f"{arm:<10}" + "".join(cells) + f"{mb:>9.1f}{100*anyv:>9.2f}")

    print(f"\nPER-RUN MECHANISM vs `{ref}` (two-sided permutation on per-run means; "
          f"exact when C(n,na)<=2e7)")
    print(f"{'metric':<24} {'arm':<10} {'delta':>9} {'p':>9} {'method':<12} "
          f"{'MDE':>9}")
    print("-" * 78)
    for key, label, scale, unit in metrics:
        xa = per_run(ref, key)
        if len(xa) < 2:
            continue
        sd = _sd(xa)
        mde = Z_ALPHA_POWER * sd * math.sqrt(2.0 / len(xa)) * scale
        for arm in arms:
            if arm == ref:
                continue
            xb = per_run(arm, key)
            res = perm_test(xa, xb)
            if res is None:
                continue
            obs, p, method = res
            signed = (_mean(xb) - _mean(xa)) * scale
            print(f"{label:<24} {arm:<10} {signed:>+9.2f} {p:>9.4f} "
                  f"{method:<12} {mde:>8.2f}{unit}")
    return 0


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