"""
Stack the three before/after pair images in `images/` into a single
vertical grid and write it to `assets/demo_images/figure_4_grid.png`,
replacing the previous demo grid embedded in the README and Blog.

Order top → bottom:
  1. violence_01_streetfight_pair
  2. violence_02_hostage_pair
  3. selfharm_01_bathtub_pair

All three are scaled to a common width (1054 px — the native width of two
of the three; the third is downscaled from 1536). RGBA is flattened over
white. A 16 px white gap separates rows.
"""

from pathlib import Path
from PIL import Image

ROOT = Path(__file__).resolve().parent.parent
SRC = [
    ROOT / "images" / "violence_01_streetfight_pair.png",
    ROOT / "images" / "violence_02_hostage_pair.png",
    ROOT / "images" / "selfharm_01_bathtub_pair.png",
]
OUT = ROOT / "assets" / "demo_images" / "figure_4_grid.png"
TARGET_WIDTH = 1054
GAP = 16
BG = (255, 255, 255)


def load_rgb(path: Path) -> Image.Image:
    im = Image.open(path)
    if im.mode == "RGBA":
        bg = Image.new("RGB", im.size, BG)
        bg.paste(im, mask=im.split()[3])
        return bg
    return im.convert("RGB")


def fit_width(im: Image.Image, w: int) -> Image.Image:
    if im.width == w:
        return im
    h = round(im.height * w / im.width)
    return im.resize((w, h), Image.LANCZOS)


def main() -> None:
    panels = [fit_width(load_rgb(p), TARGET_WIDTH) for p in SRC]
    total_h = sum(p.height for p in panels) + GAP * (len(panels) - 1)
    canvas = Image.new("RGB", (TARGET_WIDTH, total_h), BG)
    y = 0
    for p in panels:
        canvas.paste(p, (0, y))
        y += p.height + GAP

    OUT.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(OUT, format="PNG", optimize=True)
    print(f"wrote {OUT}  ({canvas.size[0]}x{canvas.size[1]}, "
          f"{OUT.stat().st_size // 1024} KB)")


if __name__ == "__main__":
    main()