"""Control baseline: does our optimized perturbation beat *random noise* at a
matched perceptual budget? Generates noise-perturbed copies of the test images
into a run dir (adv/<stem>.png) so scripts/eval_frontier_adv.py can score them
against the frontier VLMs exactly like a real attack run.

Two budgets:
  eps6      uniform L-inf noise at eps/255 (same L-inf budget the attack got)
  ssimmatch per-image Gaussian noise, sigma bisected to hit a target SSIM
            (default: the per-image SSIM our attack achieved, from a ref run)

  python scripts/make_noise_baseline.py --manifest examples/testset.csv \
      --images examples/testset --mode eps6 --eps 6 --out runs/noise_eps6
  python scripts/make_noise_baseline.py --manifest examples/testset.csv \
      --images examples/testset --mode ssimmatch --ref-run runs/p2_v2_e6_plain \
      --out runs/noise_ssimmatch
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
from PIL import Image
from skimage.metrics import structural_similarity as ssim_fn


def load(p: Path) -> np.ndarray:
    return np.asarray(Image.open(p).convert("RGB"), dtype=np.float32) / 255.0


def save(a: np.ndarray, p: Path) -> None:
    Image.fromarray((a.clip(0, 1) * 255).round().astype(np.uint8)).save(p)


def ssim_of(a: np.ndarray, b: np.ndarray) -> float:
    return float(ssim_fn(a, b, channel_axis=2, data_range=1.0))


def uniform_eps(rng: np.random.Generator, img: np.ndarray, eps255: float) -> np.ndarray:
    e = eps255 / 255.0
    n = rng.uniform(-e, e, size=img.shape).astype(np.float32)
    return (img + n).clip(0, 1)


def gaussian_to_ssim(rng: np.random.Generator, img: np.ndarray, target_ssim: float) -> tuple[np.ndarray, float]:
    """Bisect Gaussian sigma so the noisy image's SSIM ~= target (lower SSIM = more noise)."""
    lo, hi = 0.0, 0.5  # sigma in [0,1] pixel units
    base_noise = rng.standard_normal(img.shape).astype(np.float32)
    best = img.copy()
    for _ in range(18):
        mid = (lo + hi) / 2
        cand = (img + base_noise * mid).clip(0, 1)
        s = ssim_of(cand, img)
        if s > target_ssim:   # too clean -> more noise
            lo = mid
        else:                 # too noisy -> less noise
            hi = mid
        best = cand
    return best, ssim_of(best, img)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--images", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--mode", choices=["eps6", "ssimmatch"], required=True)
    ap.add_argument("--eps", type=float, default=6.0)
    ap.add_argument("--ref-run", default="", help="run dir with results.json for per-image target SSIM")
    ap.add_argument("--target-ssim", type=float, default=0.94)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    out = Path(args.out)
    (out / "adv").mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(args.seed)
    imgdir = Path(args.images)

    ref_ssim: dict[str, float] = {}
    if args.ref_run:
        ref = json.loads((Path(args.ref_run) / "results.json").read_text())
        for r in ref:
            if "stealth" in r:
                ref_ssim[Path(r["image"]).stem] = r["stealth"]["ssim"]

    rows = []
    for line in Path(args.manifest).read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#"):
            p, _ = line.split(",", 1)
            rows.append(Path(p).name)

    report = []
    for fname in rows:
        fp = imgdir / fname
        if not fp.exists():
            continue
        stem = fp.stem
        img = load(fp)
        if args.mode == "eps6":
            noisy = uniform_eps(rng, img, args.eps)
            achieved = None
        else:
            tgt = ref_ssim.get(stem, args.target_ssim)
            noisy, achieved = gaussian_to_ssim(rng, img, tgt)
        save(noisy, out / "adv" / f"{stem}.png")
        report.append({
            "image": fname, "ssim": round(ssim_of(noisy, img), 4),
            "target_ssim": round(ref_ssim.get(stem, args.target_ssim), 4) if args.mode == "ssimmatch" else None,
        })

    (out / "baseline_meta.json").write_text(json.dumps(
        {"mode": args.mode, "eps": args.eps, "images": report}, indent=2))
    ss = [r["ssim"] for r in report]
    print(f"wrote {len(report)} noisy images to {out}/adv  mean SSIM={sum(ss)/len(ss):.4f}")


if __name__ == "__main__":
    main()