scripts/aggregate_p2.py

"""Aggregate the P2 sweep into a stealth-vs-transfer Pareto table.

For each runs/p2_* cell: mean stealth (PSNR/SSIM/ΔE/LPIPS from results.json), open
held-out flip rate (from results.json), and frontier flip rate (from frontier_eval.json
if present). Prints a table sorted by tower/eps/perceptual so v1(ablation) vs v2 is
directly comparable at each stealth point.

  python scripts/aggregate_p2.py --runs runs            # after pulling cells to Mac
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def cell_stats(cell: Path) -> dict | None:
    rj = cell / "results.json"
    if not rj.exists():
        return None
    res = json.loads(rj.read_text())
    if not res:
        return None

    def flip_rate(grp):
        f = n = 0
        for r in res:
            for d in r.get(grp, {}).values():
                n += 1
                f += (d["adv_jpeg"] > 0 and d["clean"] < 0)
        return (f / n if n else 0.0, f, n)

    st_keys = [k for k in ("psnr", "ssim", "deltaE_p95", "lpips")
               if "stealth" in res[0] and k in res[0]["stealth"]]
    stealth = {k: sum(r["stealth"][k] for r in res if "stealth" in r) / len(res)
               for k in st_keys}

    hr, hf, hn = flip_rate("heldout")
    tr, tf, tn = flip_rate("train")

    out = {"cell": cell.name, "images": len(res),
           "train_flip": tr, "heldout_flip": hr,
           "train_fn": f"{tf}/{tn}", "heldout_fn": f"{hf}/{hn}", **stealth}

    out["front_flip"] = 0.0
    out["_front_f"] = out["_front_n"] = 0
    fe = cell / "frontier_eval.json"
    if fe.exists():
        fdata = json.loads(fe.read_text())
        tal: dict[str, list[int]] = {}
        for r in fdata:
            for m, d in r["models"].items():
                tal.setdefault(m, [0, 0])
                tal[m][0] += bool(d["flip"])
                tal[m][1] += 1
        tf = tn = 0
        for m, (f, n) in tal.items():
            short = m.split("/")[-1]
            out[f"front_{short}"] = f"{f}/{n}"
            tf += f
            tn += n
        out["_front_f"], out["_front_n"] = tf, tn
        out["front_flip"] = tf / tn if tn else 0.0
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--runs", default="runs")
    ap.add_argument("--glob", default="p2_*")
    args = ap.parse_args()

    cells = sorted(Path(args.runs).glob(args.glob))
    rows = [s for c in cells if (s := cell_stats(c))]
    if not rows:
        print("no cells found")
        return

    # rank by frontier transfer (primary), then held-out open transfer
    rows.sort(key=lambda r: (r.get("front_flip", 0.0), r["heldout_flip"]), reverse=True)
    for i, r in enumerate(rows, 1):
        r["rank"] = i

    front_named = sorted({k for r in rows for k in r
                          if k.startswith("front_") and k != "front_flip"})
    cols = (["rank", "cell", "images", "heldout_fn", "heldout_flip"]
            + front_named + ["front_flip", "psnr", "ssim", "deltaE_p95", "lpips"])

    def fmt(v):
        return f"{v:.3f}" if isinstance(v, float) else str(v)

    w = {c: max(len(c), *(len(fmt(r.get(c, ""))) for r in rows)) for c in cols}
    print("  ".join(c.ljust(w[c]) for c in cols))
    print("  ".join("-" * w[c] for c in cols))
    for r in rows:
        print("  ".join(fmt(r.get(c, "")).ljust(w[c]) for c in cols))

    # SSIM-ranked view (stealth), highest first
    print("\n--- ranked by SSIM (stealth) ---")
    for i, r in enumerate(sorted(rows, key=lambda x: x.get("ssim", 0), reverse=True), 1):
        print(f"  {i}. {r['cell']:<24} ssim={r.get('ssim',0):.4f}  "
              f"lpips={r.get('lpips',0):.4f}  psnr={r.get('psnr',0):.2f}  "
              f"heldout={r['heldout_fn']}  frontier={r.get('_front_f',0)}/{r.get('_front_n',0)}")


if __name__ == "__main__":
    main()