scripts/smoke_whitebox.py

"""White-box pivot smoke battery — cheap tests, runnable in well under an hour.

Batteries:
  A  Encoder strength ladder (free): does faint->loud text monotonically move the
     decoy margin on CLIP + SigLIP2? (tests the fix for the saturation bug)
  C  Gradient path (free): nonzero region gradient + a single ascent step raises
     the margin, on both encoders.
  B  Proxy validity (surrogate calls): does encoder delta-margin predict a LOCAL
     VLM (qwen/gemma) label flip? Spearman + flip-rate by strength.
  D  Local transfer control (surrogate calls): does a far/high-margin decoy overlay
     flip local VLMs more than a near/control decoy at the same strength?
  E  Tiny frontier probe (small API spend): do the loud overlays flip gpt-5.5 /
     gemini-3.5-flash vs the clean image, on a handful of images?

Writes research/smoke_results.json and prints a summary.
"""

from __future__ import annotations

import json
import sys
import time
from pathlib import Path

from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from veil_pgd.config import get_settings, stealth_preset  # noqa: E402
from veil_pgd.fitness.embed import Embedder  # noqa: E402
from veil_pgd.fitness.semantic import embedding_distance  # noqa: E402
from veil_pgd.render import render  # noqa: E402
from veil_pgd.robustness import scraper_sim  # noqa: E402
from veil_pgd.stealth import evaluate_stealth  # noqa: E402
from veil_pgd.targets.base import LabelPrompt  # noqa: E402
from veil_pgd.targets.registry import Registry  # noqa: E402
from veil_pgd.targets.whitebox import WhiteBoxClient  # noqa: E402
from veil_pgd.types import RenderSpec  # noqa: E402

ENCODERS = ["openclip:ViT-B-32", "siglip2-base"]
DECOY_POOL = [
    "banana", "jellyfish", "volcano", "umbrella", "bonfire", "tarantula",
    "lighthouse", "waterfall", "cactus", "chandelier", "accordion", "igloo",
    "peacock", "typewriter", "pineapple", "seahorse", "windmill", "hammock",
    "lantern", "pretzel", "strawberry", "octopus",
]
FLIP_TAU = 0.5


def strength_specs(decoy: str) -> dict[str, RenderSpec]:
    # font_px is @1024 reference; scale_px shrinks it to the image short side, so
    # on ~320px test images these map to actual ~9 / 15 / 27 px. "loud" is meant
    # to be comfortably readable by a VLM; "faint" sits near the readability floor.
    return {
        "faint": RenderSpec(text=decoy, font_px=32, alpha=0.25,
                            color_strategy="bg_mean_offset", brightness_offset=20,
                            position="bottom_middle"),
        "mid": RenderSpec(text=decoy, font_px=52, alpha=0.45,
                         color_strategy="bg_mean_offset", brightness_offset=30,
                         position="bottom_middle"),
        "loud": RenderSpec(text=decoy, font_px=88, alpha=0.80,
                          color_strategy="fixed", fixed_rgb=(10, 10, 10),
                          position="bottom_middle", repetition=2),
    }


def log(m: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)


def read_rows(path: str, limit: int | None = None) -> list[tuple[str, str]]:
    rows = []
    for line in Path(path).read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        p, t = line.split(",", 1)
        rows.append((p.strip(), t.strip()))
    if limit:
        # spread across classes: take every Nth
        step = max(1, len(rows) // limit)
        rows = rows[::step][:limit]
    return rows


def spearman(xs: list[float], ys: list[float]) -> float:
    n = len(xs)
    if n < 3:
        return float("nan")

    def ranks(v):
        order = sorted(range(len(v)), key=lambda i: v[i])
        r = [0.0] * len(v)
        for rank, i in enumerate(order):
            r[i] = rank
        return r

    rx, ry = ranks(xs), ranks(ys)
    mx, my = sum(rx) / n, sum(ry) / n
    num = sum((rx[i] - mx) * (ry[i] - my) for i in range(n))
    dx = sum((rx[i] - mx) ** 2 for i in range(n)) ** 0.5
    dy = sum((ry[i] - my) ** 2 for i in range(n)) ** 0.5
    return num / (dx * dy) if dx and dy else float("nan")


def pick_decoys(embedder: Embedder, truth: str) -> tuple[str, str]:
    """far decoy = pool word most distant from truth; near = least distant."""
    cands = [w for w in DECOY_POOL if w != truth]
    dists = [(w, embedding_distance(embedder, w, truth)) for w in cands]
    dists.sort(key=lambda kv: kv[1])
    return dists[-1][0], dists[0][0]  # (far, near)


def main():
    s = get_settings()
    reg = Registry(s)
    wb = WhiteBoxClient(s.klaus3_vision_service_url)
    embedder = Embedder(reg.embeddings(), s.klaus3_vision_service_url)
    prompt = LabelPrompt()
    strict = stealth_preset("strict")

    # Use whichever local surrogates are actually reachable (gemma may be paused
    # to free VRAM for the encoders).
    import httpx
    surrogates = []
    for name, url in [("qwen-3.5-4b", s.klaus3_qwen_base_url),
                      ("gemma-4-4b", s.klaus3_gemma4b_base_url)]:
        try:
            httpx.get(url.rstrip("/") + "/models", timeout=3.0)
            surrogates.append(reg.surrogate(name))
        except Exception:
            log(f"surrogate {name} unreachable, skipping")
    log(f"active surrogates: {[m.name for m in surrogates]}")

    log(f"vit health: {wb.health()['loaded']} available")
    for e in ENCODERS:
        wb.load(e)
    log(f"encoders loaded: {[e for e in ENCODERS]}")

    rows_free = read_rows("examples/testset.csv", limit=20)
    rows_surr = rows_free[:12]
    rows_transfer = rows_free[:8]
    rows_frontier = rows_free[:6]

    results: dict = {"config": {"encoders": ENCODERS, "flip_tau": FLIP_TAU},
                     "A_strength_ladder": [], "C_gradient": [],
                     "B_proxy": [], "D_transfer": [], "E_frontier": []}

    # decoy cache
    decoys: dict[str, tuple[str, str]] = {}
    for _, truth in rows_free:
        if truth not in decoys:
            decoys[truth] = pick_decoys(embedder, truth)

    # ---- Battery A: strength ladder on encoders (free) ----
    log("== Battery A: encoder strength ladder ==")
    for path, truth in rows_free:
        img = Image.open(path).convert("RGB")
        far, _ = decoys[truth]
        specs = strength_specs(far)
        row = {"image": Path(path).name, "truth": truth, "decoy": far, "by_encoder": {}}
        for enc in ENCODERS:
            deltas = {}
            for level, spec in specs.items():
                cand = scraper_sim(render(img, spec))
                r = wb.score(cand, truth, far, model_id=enc, clean=img)
                deltas[level] = r["delta_margin"]
            mono = deltas["faint"] <= deltas["mid"] <= deltas["loud"]
            row["by_encoder"][enc] = {"delta": deltas, "monotonic_up": mono}
        results["A_strength_ladder"].append(row)
    log(f"  A done: {len(results['A_strength_ladder'])} images")

    # ---- Battery C: gradient path (free) ----
    log("== Battery C: gradient path ==")
    for path, truth in rows_free:
        img = Image.open(path).convert("RGB")
        far, _ = decoys[truth]
        W, H = img.size
        region = [0, int(H * 0.78), W, H]
        cand = scraper_sim(render(img, strength_specs(far)["mid"]))
        row = {"image": Path(path).name, "truth": truth, "decoy": far, "by_encoder": {}}
        for enc in ENCODERS:
            g = wb.grad_region(cand, truth, far, model_id=enc, region=region)
            row["by_encoder"][enc] = {
                "grad_l2": g["grad_l2"], "grad_l2_region": g["grad_l2_region"],
                "margin_increased": g["margin_increased"]}
        results["C_gradient"].append(row)
    log(f"  C done: {len(results['C_gradient'])} images")

    # ---- Battery B: does encoder margin predict local VLM flip? ----
    log("== Battery B: proxy validity (surrogate flips) ==")
    b_margin_clip, b_surr_dist, b_flip = [], [], []
    for path, truth in rows_surr:
        img = Image.open(path).convert("RGB")
        far, _ = decoys[truth]
        for level, spec in strength_specs(far).items():
            cand = scraper_sim(render(img, spec))
            clip_margin = wb.score(cand, truth, far, model_id="openclip:ViT-B-32",
                                   clean=img)["delta_margin"]
            dists = []
            preds = {}
            for m in surrogates:
                res = m.label(cand, prompt)
                preds[m.name] = res.parsed_label
                dists.append(embedding_distance(embedder, res.parsed_label, truth))
            surr_dist = sum(dists) / len(dists)
            flipped = surr_dist >= FLIP_TAU
            b_margin_clip.append(clip_margin)
            b_surr_dist.append(surr_dist)
            b_flip.append(1.0 if flipped else 0.0)
            results["B_proxy"].append({
                "image": Path(path).name, "truth": truth, "decoy": far,
                "level": level, "clip_delta_margin": clip_margin,
                "surrogate_dist": surr_dist, "flipped": flipped, "preds": preds})
    results["B_proxy_summary"] = {
        "spearman_margin_vs_surrdist": spearman(b_margin_clip, b_surr_dist),
        "spearman_margin_vs_flip": spearman(b_margin_clip, b_flip),
        "flip_rate": sum(b_flip) / len(b_flip) if b_flip else 0.0,
        "n": len(b_flip)}
    log(f"  B done: {results['B_proxy_summary']}")

    # ---- Battery D: far/high-margin decoy vs near control ----
    log("== Battery D: local transfer control ==")
    d_far_flips, d_near_flips = 0, 0
    for path, truth in rows_transfer:
        img = Image.open(path).convert("RGB")
        far, near = decoys[truth]
        row = {"image": Path(path).name, "truth": truth, "far": far, "near": near}
        for tag, decoy in [("far", far), ("near", near)]:
            cand = scraper_sim(render(img, strength_specs(decoy)["loud"]))
            dists = []
            for m in surrogates:
                res = m.label(cand, prompt)
                dists.append(embedding_distance(embedder, res.parsed_label, truth))
            sd = sum(dists) / len(dists)
            row[f"{tag}_dist"] = sd
            row[f"{tag}_flip"] = sd >= FLIP_TAU
            if tag == "far" and sd >= FLIP_TAU:
                d_far_flips += 1
            if tag == "near" and sd >= FLIP_TAU:
                d_near_flips += 1
        results["D_transfer"].append(row)
    results["D_transfer_summary"] = {
        "n": len(rows_transfer), "far_flips": d_far_flips, "near_flips": d_near_flips}
    log(f"  D done: {results['D_transfer_summary']}")

    # ---- Battery E: tiny frontier probe (paid) ----
    log("== Battery E: frontier probe (gpt-5.5 + gemini) ==")
    blackbox = reg.all_blackbox()
    for path, truth in rows_frontier:
        img = Image.open(path).convert("RGB")
        far, _ = decoys[truth]
        loud = scraper_sim(render(img, strength_specs(far)["loud"]))
        # stealth of the loud overlay (no lpips call, just psnr/ssim/dE)
        st = evaluate_stealth(img, render(img, strength_specs(far)["loud"]),
                              thresholds=strict, lpips_fn=None)
        row = {"image": Path(path).name, "truth": truth, "decoy": far,
               "loud_stealth": {"psnr": st.psnr, "ssim": st.ssim,
                                "delta_e_p95": st.delta_e_p95, "passed": st.passed},
               "by_model": {}}
        for m in blackbox:
            clean_pred = m.label(img, prompt).parsed_label
            adv_pred = m.label(loud, prompt).parsed_label
            clean_d = embedding_distance(embedder, clean_pred, truth)
            adv_d = embedding_distance(embedder, adv_pred, truth)
            row["by_model"][m.name] = {
                "clean_pred": clean_pred, "adv_pred": adv_pred,
                "clean_dist": clean_d, "adv_dist": adv_d,
                "flipped": adv_d >= FLIP_TAU and clean_d < FLIP_TAU}
        results["E_frontier"].append(row)
        log(f"  {Path(path).name} truth={truth} decoy={far}: "
            + ", ".join(f"{k.split('/')[-1]}={'FLIP' if v['flipped'] else 'no'}"
                        f"({v['adv_pred']!r})" for k, v in row['by_model'].items()))

    out = Path("research/smoke_results.json")
    out.write_text(json.dumps(results, indent=2))
    log(f"wrote {out}")

    # ---- summary ----
    print("\n================ SMOKE SUMMARY ================")
    for enc in ENCODERS:
        mono = [r["by_encoder"][enc]["monotonic_up"] for r in results["A_strength_ladder"]]
        avg_loud = sum(r["by_encoder"][enc]["delta"]["loud"]
                       for r in results["A_strength_ladder"]) / len(mono)
        print(f"A {enc}: monotonic-up {sum(mono)}/{len(mono)}; "
              f"mean loud delta_margin={avg_loud:+.3f}")
    for enc in ENCODERS:
        inc = [r["by_encoder"][enc]["margin_increased"] for r in results["C_gradient"]]
        print(f"C {enc}: grad step raised margin {sum(inc)}/{len(inc)}")
    bs = results["B_proxy_summary"]
    print(f"B proxy: spearman(margin,surr_dist)={bs['spearman_margin_vs_surrdist']:.2f} "
          f"spearman(margin,flip)={bs['spearman_margin_vs_flip']:.2f} "
          f"local flip_rate={bs['flip_rate']:.0%} (n={bs['n']})")
    ds = results["D_transfer_summary"]
    print(f"D transfer: far-decoy flips {ds['far_flips']}/{ds['n']} vs "
          f"near-control {ds['near_flips']}/{ds['n']}")
    e_flips = {}
    for r in results["E_frontier"]:
        for k, v in r["by_model"].items():
            e_flips.setdefault(k, [0, 0])
            e_flips[k][0] += 1 if v["flipped"] else 0
            e_flips[k][1] += 1
    e_pass = sum(1 for r in results["E_frontier"] if r["loud_stealth"]["passed"])
    print(f"E frontier (loud overlay): "
          + "; ".join(f"{k.split('/')[-1]} {v[0]}/{v[1]} flipped" for k, v in e_flips.items())
          + f"; loud passed STRICT stealth {e_pass}/{len(results['E_frontier'])}")
    print("===============================================")

    reg.close()
    wb.close()


if __name__ == "__main__":
    main()