"""Rebuild the README figures from a completed run (clean vs protected).

  python scripts/make_figures.py --run runs/v021_base

Writes:
  assets/compare_clean_vs_adv.png  a few clean|protected pairs, high-res
  assets/grid_before_after_16.png  16 images, clean over protected, 4x4
"""
from __future__ import annotations

import argparse
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

IMAGES = Path("examples/testset60")
ASSETS = Path("assets")
CELL = 320
PAD = 14
LABEL_H = 26


def load_manifest() -> list[tuple[str, str]]:
    rows = []
    for line in Path("examples/testset60.csv").read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#"):
            p, t = line.split(",", 1)
            rows.append((Path(p).name, t.strip()))
    return rows


def font(size: int) -> ImageFont.FreeTypeFont:
    for p in ("/System/Library/Fonts/Supplemental/Arial.ttf",
              "/System/Library/Fonts/Helvetica.ttc",
              "/Library/Fonts/Arial.ttf"):
        if Path(p).exists():
            return ImageFont.truetype(p, size)
    return ImageFont.load_default()


def sq(img: Image.Image, n: int = CELL) -> Image.Image:
    return img.convert("RGB").resize((n, n), Image.LANCZOS)


def labeled(img: Image.Image, text: str, f) -> Image.Image:
    canvas = Image.new("RGB", (img.width, img.height + LABEL_H), "white")
    canvas.paste(img, (0, LABEL_H))
    d = ImageDraw.Draw(canvas)
    d.text((4, 4), text, fill="black", font=f)
    return canvas


def compare(run: Path, rows: list[tuple[str, str]], picks: list[str]) -> None:
    f = font(18)
    fh = font(20)
    items = [(fn, t) for fn, t in rows if fn in picks]
    w = 2 * CELL + 3 * PAD
    h = len(items) * (CELL + LABEL_H + PAD) + PAD + 30
    canvas = Image.new("RGB", (w, h), "white")
    d = ImageDraw.Draw(canvas)
    d.text((PAD, 6), "Clean", fill="black", font=fh)
    d.text((2 * PAD + CELL, 6), "Protected (VEIL-PGD)", fill="black", font=fh)
    y = 34
    for fn, truth in items:
        stem = Path(fn).stem
        clean = sq(Image.open(IMAGES / fn))
        adv = sq(Image.open(run / "adv" / f"{stem}.png"))
        canvas.paste(labeled(clean, truth, f), (PAD, y))
        canvas.paste(labeled(adv, truth, f), (2 * PAD + CELL, y))
        y += CELL + LABEL_H + PAD
    ASSETS.mkdir(exist_ok=True)
    canvas.save(ASSETS / "compare_clean_vs_adv.png")
    print("wrote assets/compare_clean_vs_adv.png", canvas.size)


def grid(run: Path, rows: list[tuple[str, str]], picks: list[str]) -> None:
    f = font(16)
    items = [(fn, t) for fn, t in rows if fn in picks][:16]
    cols = 4
    cellw = CELL
    cellh = 2 * CELL + LABEL_H
    n = len(items)
    grows = (n + cols - 1) // cols
    w = cols * cellw + (cols + 1) * PAD
    h = grows * (cellh + PAD) + PAD
    canvas = Image.new("RGB", (w, h), "white")
    for i, (fn, truth) in enumerate(items):
        r, c = divmod(i, cols)
        x = PAD + c * (cellw + PAD)
        y = PAD + r * (cellh + PAD)
        stem = Path(fn).stem
        clean = sq(Image.open(IMAGES / fn))
        adv = sq(Image.open(run / "adv" / f"{stem}.png"))
        block = Image.new("RGB", (cellw, cellh), "white")
        d = ImageDraw.Draw(block)
        d.text((4, 4), truth, fill="black", font=f)
        block.paste(clean, (0, LABEL_H))
        block.paste(adv, (0, LABEL_H + CELL))
        canvas.paste(block, (x, y))
    ASSETS.mkdir(exist_ok=True)
    canvas.save(ASSETS / "grid_before_after_16.png")
    print("wrote assets/grid_before_after_16.png", canvas.size)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", default="runs/v021_base")
    args = ap.parse_args()
    run = Path(args.run)
    rows = load_manifest()
    have = [fn for fn, _ in rows if (run / "adv" / f"{Path(fn).stem}.png").exists()]
    # one image per class for the small comparison
    seen, comp_picks = set(), []
    for fn, t in rows:
        if fn in have and t not in seen:
            seen.add(t)
            comp_picks.append(fn)
    compare(run, rows, comp_picks[:4])
    # 16 images spread across classes for the grid
    grid_picks = []
    by_class: dict[str, list[str]] = {}
    for fn, t in rows:
        if fn in have:
            by_class.setdefault(t, []).append(fn)
    rr = 0
    while len(grid_picks) < 16 and any(len(v) > rr for v in by_class.values()):
        for t in by_class:
            if len(by_class[t]) > rr and len(grid_picks) < 16:
                grid_picks.append(by_class[t][rr])
        rr += 1
    grid(run, rows, grid_picks)


if __name__ == "__main__":
    main()