#!/usr/bin/env python3
"""Task A/B harness — how many enemy shots is STRAFE blind to? (job j133).

The shipped detector (`strafe.nim` `detectFires`) classifies the enemy energy
delta between consecutive readings with ONE window:

    if drop >= 0.09 and drop <= 3.01: spawn one wave with power = drop

This script replays that detector, and the j133 fixed detector, against the
recorded per-tick enemy energy trace of a battle corpus and compares each
created wave against the TRUE fire events (`*.events.jsonl`, `type == "fire"`).
It is the honest size of the "missed fire" problem, offline.

Server facts that make the window wrong (verified in the SERVER source):

  * `rules.kt BULLET_HIT_ENERGY_GAIN_FACTOR = 3`: the SHOOTER gains `3*power`
    when its bullet hits the victim.  When an enemy bullet hits us and the enemy
    fires in the SAME tick, the gain cancels the fire cost and the delta reads
    as "no fire".
  * our own bullet damaging the enemy in the same tick adds `damage` to the
    delta and can push it past 3.01 -> REJECTED, no heat for the enemy's shot.

The FIXED detector corrects the delta by exactly those two observable amounts
(`onHitByBullet` / `onBulletHit` supply them live) and SPLITS a still-too-large
delta instead of dropping it.

Run:
  python3 common_libs/tests/measure_strafe_fire_catch.py \
      --corpus /tmp/tfil_ab2/out \
      --report common_libs/tests/fixtures/strafe_fire_catch_report.txt

Corpus rows are the perfect-information trajectory (`ex,ey,ee,...`); the live
bot's WorldState carries the SAME per-tick energy because the 1v1 lock radar
scans every tick (see `--worldstate` to check that on live-recorded captures).
"""
from __future__ import annotations

import argparse
import collections
import json
import math
import os
import sys

# Corpus artifact: the worldstate row `t` holds the energy AFTER the events whose
# round-tick is `t+1` (the events side-car and the capture are offset by one).
# The LIVE bot gets the event and the reading in the SAME turn, so this shift is
# a corpus artifact only.
CORPUS_EVENT_SHIFT = 1

# Per-mover detector windows, as wired in j134.  Every mover now calls the SAME
# shared `movement_harness/fire_tracker.nim`; only the window it supplies differs
# (`strafe` is the j133 reference).  The correction and the over-cap SPLIT are
# shared, so the FIXED catch rate must be identical for all of them; the SHIPPED
# rate can differ slightly where the windows differ.
MOVER_WINDOWS = {
    "tfil":    (0.09, 3.01),
    "ring":    (0.09, 3.01),
    "strafe":  (0.09, 3.01),
    "learned": (0.09, 3.01),
    "surf":    (0.10, 3.00),
}


def load_run(path: str):
    rows, events = [], []
    with open(path + ".jsonl") as f:
        for line in f:
            d = json.loads(line)
            if "tick" in d:
                rows.append(d)
    with open(path + ".events.jsonl") as f:
        for line in f:
            line = line.strip()
            if line:
                events.append(json.loads(line))
    with open(path + ".jsonl.rounds.json") as f:
        rounds = json.load(f)["rounds"]
    return rows, events, rounds


def detect(rows, events, rounds, fix: bool, lo: float = 0.09, hi: float = 3.01):
    """Replay the detector; return {global_tick: [powers]} of created waves."""
    starts = {r["round"]: r["startTick"] for r in rounds}
    by_global = {r["startTick"] + k: (r["round"], k)
                 for r in rounds for k in range(r["count"])}
    energy = {r["tick"]: r["ee"] for r in rows}
    pos = {r["tick"]: r for r in rows}

    # Enemy owner = the one whose fire POSITION matches (`ex`,`ey`).
    votes = collections.Counter()
    for e in events:
        if e["type"] != "fire":
            continue
        gt = starts.get(e["round"])
        if gt is None:
            continue
        for off in (-1, -2, 0):
            d = pos.get(gt + e["tick"] + off)
            if d is None:
                continue
            de = math.hypot(d["ex"] - e["x"], d["ey"] - e["y"])
            ds = math.hypot(d["sx"] - e["x"], d["sy"] - e["y"])
            votes[(e["owner"], "e" if de < 2.0 else "s")] += 1
            break
    eowner = next((o for (o, w), _ in votes.most_common() if w == "e"), None)
    if eowner is None:
        return {}, energy, eowner

    bullet_power = {(e["round"], e["bullet"]): e["power"]
                    for e in events if e["type"] == "fire" and e["owner"] == eowner}
    bonus = collections.defaultdict(float)   # 3*power of enemy hits on us
    dealt = collections.defaultdict(float)   # our damage to the enemy
    for e in events:
        if e["type"] != "hit":
            continue
        if e.get("owner") == eowner and e.get("victim") != eowner:
            p = bullet_power.get((e["round"], e["bullet"]))
            if p is not None:
                bonus[(e["round"], e["tick"])] += 3.0 * p
        if e.get("victim") == eowner and e.get("owner") != eowner:
            dealt[(e["round"], e["tick"])] += e.get("damage", 0.0)

    waves = collections.defaultdict(list)
    prev = 100.0
    for t in sorted(energy):
        raw = prev - energy[t]
        prev = energy[t]
        rn, rt = by_global.get(t, (0, t))
        drop = raw
        if fix:
            drop = raw + bonus.get((rn, rt + CORPUS_EVENT_SHIFT), 0.0) \
                      - dealt.get((rn, rt + CORPUS_EVENT_SHIFT), 0.0)
        if fix and drop > hi:
            n = int(math.ceil(drop / 3.0))
            for _ in range(n):
                waves[t].append(drop / n)
        elif lo <= drop <= hi:
            waves[t].append(drop)
    return waves, energy, eowner


def fires_of(events, rounds, eowner):
    starts = {r["round"]: r["startTick"] for r in rounds}
    out = []
    for e in events:
        if e["type"] != "fire" or e["owner"] != eowner:
            continue
        gt = starts.get(e["round"])
        if gt is not None:
            out.append((gt + e["tick"], e["power"]))
    return out


def score(rows, events, rounds, fix: bool, lo: float = 0.09, hi: float = 3.01):
    waves, energy, eowner = detect(rows, events, rounds, fix, lo, hi)
    starts = {r["round"]: r["startTick"] for r in rounds}
    by_global = {r["startTick"] + k: (r["round"], k)
                 for r in rounds for k in range(r["count"])}
    total = caught = 0
    lat = collections.Counter()
    cats = collections.Counter()
    powers = []
    for gt, power in fires_of(events, rounds, eowner):
        dt = gt - 1                       # row holding the energy change
        if dt - 1 not in energy or dt not in energy:
            continue
        total += 1
        if waves.get(dt):
            caught += 1
            lat[0] += 1
            powers.append(waves[dt][0] - power)
            continue
        rn, rt = by_global.get(dt, (0, dt))
        raw = energy[dt - 1] - energy[dt]
        drop = raw
        if fix:
            drop = raw  # classification below only needs the miss kind
        if raw > hi:
            cats["reject_high"] += 1
        else:
            cats["reject_low"] += 1
        nxt = -1
        for k in range(1, 6):
            if waves.get(dt + k):
                nxt = k
                break
        lat[nxt] += 1
    return total, caught, lat, cats


def scan_interval(files):
    incs = collections.Counter()
    ticks = 0
    for fn in files:
        rows = []
        with open(fn) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                d = json.loads(line)
                if "tick" in d:
                    rows.append(d)
        for i in range(1, len(rows)):
            if "lst" in rows[i] and "lst" in rows[i - 1]:
                incs[rows[i]["lst"] - rows[i - 1]["lst"]] += 1
                ticks += 1
    return incs, ticks


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--corpus", default="/tmp/tfil_ab2/out")
    ap.add_argument("--report", default="")
    ap.add_argument("--worldstate", nargs="*", default=[],
                    help="live-recorded WorldState captures to measure the scan interval")
    args = ap.parse_args()

    arms = sorted(d for d in os.listdir(args.corpus)
                  if os.path.isdir(os.path.join(args.corpus, d)))
    runs = [p for arm in arms for p in
            (os.path.join(args.corpus, arm, f"run{r}") for r in range(1, 15))
            if os.path.exists(p + ".jsonl") and os.path.exists(p + ".events.jsonl")
            and os.path.exists(p + ".jsonl.rounds.json")]

    lines = []
    def out(s=""):
        print(s)
        lines.append(s)

    out(f"corpus: {args.corpus}  ({len(runs)} runs, {len(arms)} arms)")
    for fix in (False, True):
        total = caught = 0
        lat = collections.Counter()
        cats = collections.Counter()
        for path in runs:
            rows, events, rounds = load_run(path)
            t, c, l, k = score(rows, events, rounds, fix)
            total += t
            caught += c
            lat.update(l)
            cats.update(k)
        tag = "FIXED  (TR_FIRE_FIX=1)" if fix else "SHIPPED (TR_FIRE_FIX=0)"
        out(f"\n== {tag} ==")
        out(f"true enemy fires: {total}")
        out(f"caught (wave on the fire's own tick): {caught}  "
            f"catch rate = {caught / total:.5f}")
        out(f"missed: {total - caught}  "
            f"(drop > 3.01 [multi-fire/contamination]: {cats['reject_high']}, "
            f"drop < 0.09 [masked by the +3*power bonus]: {cats['reject_low']})")
        out("latency (ticks after the fire's own tick): " +
            ", ".join(f"{k}={v}" for k, v in sorted(lat.items())))

    if args.worldstate:
        incs, n = scan_interval(args.worldstate)
        out(f"\n== live scan interval ({n} consecutive readings over "
            f"{len(args.worldstate)} captures) ==")
        total = sum(incs.values())
        out("lst increment histogram: " + ", ".join(
            f"{k}:{v} ({100.0 * v / total:.2f}%)" for k, v in sorted(incs.items())))

    # ── per-mover table (j134) ──────────────────────────────────────────────
    # Every mover now calls the shared `movement_harness/fire_tracker.nim`; only
    # the window differs.  This is the propagation proof: the SHIPPED path leaves
    # each mover blind to the same 1.11%, the FIXED path is 100% for all.
    out("\n== per-mover catch rate (shared fire_tracker, j134) ==")
    out(f"{'mover':<8} {'window':<14} {'shipped':>9} {'fixed':>9} "
        f"{'blind_before':>13} {'blind_after':>12}")
    for name, (wlo, whi) in MOVER_WINDOWS.items():
        per = {}
        for fix in (False, True):
            total = caught = 0
            for path in runs:
                rows, events, rounds = load_run(path)
                t, c, _, _ = score(rows, events, rounds, fix, wlo, whi)
                total += t
                caught += c
            per[fix] = (total, caught)
        total, caught0 = per[False]
        _, caught1 = per[True]
        out(f"{name:<8} {f'{wlo}-{whi}':<14} {caught0 / total:>9.5f} "
            f"{caught1 / total:>9.5f} {total - caught0:>13} {total - caught1:>12}")

    if args.report:
        os.makedirs(os.path.dirname(args.report) or ".", exist_ok=True)
        with open(args.report, "w") as f:
            f.write("\n".join(lines) + "\n")
        print(f"\nwrote {args.report}")
    return 0


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