"""Aggregate v0.2 encoder-margin results with bootstrap 95% CIs, and (optionally)
the paired M4 on-vs-off delta. Operates on per-image results.json produced by
ensemble.run_attack.

Flip (per encoder, per image): margin crosses 0 after JPEG (clean<0 and adv_jpeg>0),
i.e. the decoy beats the truth on the recompressed adversarial image. We report the
per-image flip FRACTION (over that image's scored encoders) so images are the unit of
resampling — matches how we'll bootstrap the frontier eval.

  python scripts/aggregate_v02.py runs/v02_e6_plain [runs/v02_e6_noM4 ...] \
      [--paired runs/v02_e6_plain runs/v02_e6_noM4]
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np


def per_image_flips(run: Path, grp: str) -> np.ndarray:
    """Return (n_images, 2) array of [flips, scored] for the given group."""
    data = json.loads((run / "results.json").read_text())
    out = []
    for r in data:
        f = n = 0
        for d in r.get(grp, {}).values():
            n += 1
            f += (d["adv_jpeg"] > 0 and d["clean"] < 0)
        out.append([f, n])
    return np.array(out, dtype=float)


def boot_rate(fn: np.ndarray, iters: int = 10000, seed: int = 0) -> tuple[float, float, float]:
    """Bootstrap mean pooled flip RATE (sum flips / sum scored) over images."""
    rng = np.random.default_rng(seed)
    idx = np.arange(len(fn))
    samples = []
    for _ in range(iters):
        b = rng.choice(idx, size=len(idx), replace=True)
        f, n = fn[b, 0].sum(), fn[b, 1].sum()
        samples.append(f / n if n else 0.0)
    s = np.array(samples)
    point = fn[:, 0].sum() / fn[:, 1].sum()
    return point, float(np.percentile(s, 2.5)), float(np.percentile(s, 97.5))


def boot_paired_delta(a: np.ndarray, b: np.ndarray, iters: int = 10000, seed: int = 0):
    """Paired bootstrap of rate(a) - rate(b) over the same images (a,b aligned)."""
    rng = np.random.default_rng(seed)
    idx = np.arange(len(a))
    d = []
    for _ in range(iters):
        s = rng.choice(idx, size=len(idx), replace=True)
        ra = a[s, 0].sum() / max(a[s, 1].sum(), 1)
        rb = b[s, 0].sum() / max(b[s, 1].sum(), 1)
        d.append(ra - rb)
    d = np.array(d)
    point = a[:, 0].sum() / a[:, 1].sum() - b[:, 0].sum() / b[:, 1].sum()
    return point, float(np.percentile(d, 2.5)), float(np.percentile(d, 97.5))


def stealth_means(run: Path) -> dict:
    data = json.loads((run / "results.json").read_text())
    keys = ["psnr", "ssim", "deltaE_p95", "lpips"]
    vals = {k: [] for k in keys}
    for r in data:
        st = r.get("stealth", {})
        for k in keys:
            if k in st:
                vals[k].append(st[k])
    return {k: (sum(v) / len(v) if v else float("nan")) for k, v in vals.items()}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("runs", nargs="+")
    ap.add_argument("--paired", nargs=2, metavar=("A", "B"),
                    help="paired bootstrap of A-B held-out + train flip-rate delta")
    args = ap.parse_args()

    print(f"{'cell':<20} {'group':<8} {'flip rate':>10}  95% CI")
    print("-" * 56)
    for rp in args.runs:
        run = Path(rp)
        st = stealth_means(run)
        for grp in ("train", "heldout"):
            fn = per_image_flips(run, grp)
            p, lo, hi = boot_rate(fn)
            print(f"{run.name:<20} {grp:<8} {p*100:>9.1f}%  [{lo*100:.1f}, {hi*100:.1f}]")
        print(f"{'':<20} stealth  ssim={st['ssim']:.3f} lpips={st['lpips']:.3f} "
              f"psnr={st['psnr']:.1f} dE95={st['deltaE_p95']:.2f}")
        print()

    if args.paired:
        A, B = Path(args.paired[0]), Path(args.paired[1])
        print(f"=== paired delta  {A.name} - {B.name}  (same 60 images) ===")
        for grp in ("train", "heldout"):
            a, b = per_image_flips(A, grp), per_image_flips(B, grp)
            p, lo, hi = boot_paired_delta(a, b)
            sig = "" if (lo <= 0 <= hi) else "  *significant (CI excludes 0)"
            print(f"  {grp:<8} Δ = {p*100:+.1f}pp  95% CI [{lo*100:+.1f}, {hi*100:+.1f}]{sig}")


if __name__ == "__main__":
    main()