from PIL import Image, ImageDraw
from pathlib import Path
import collections, re

SRC = Path(r"C:\allgifted\mathapi11v2\docs\past year questions\Only Images")
OUT = Path(r"C:\allgifted\mathapi11v2\docs\past year questions\_composites")
OUT.mkdir(parents=True, exist_ok=True)
for old in OUT.glob("*.png"): old.unlink()

MAXW, GAP, MARGIN, BORDER = 360, 14, 16, 3
D0ACAC = (208, 172, 172)

def idx(f):
    m = re.search(r"_illus_(\d+)", f.name)
    return int(m.group(1)) if m else 0

groups = collections.defaultdict(list)
for f in SRC.glob("*.png"):
    groups[f.name.split("_illus_")[0]].append(f)

print(f"{sum(len(v) for v in groups.values())} crops in {len(groups)} papers", flush=True)
for stem, files in sorted(groups.items()):
    files = sorted(files, key=idx)
    imgs = []
    for f in files:
        im = Image.open(f).convert("RGB")
        if im.width > MAXW:
            im = im.resize((MAXW, int(im.height * MAXW / im.width)), Image.LANCZOS)
        imgs.append(im)
    W = max(i.width for i in imgs) + MARGIN * 2
    H = sum(i.height for i in imgs) + GAP * (len(imgs) - 1) + MARGIN * 2
    canvas = Image.new("RGB", (W, H), (255, 255, 255))
    dr = ImageDraw.Draw(canvas)
    y = MARGIN
    for im in imgs:
        x = (W - im.width) // 2
        canvas.paste(im, (x, y))
        dr.rectangle([x-2, y-2, x+im.width+1, y+im.height+1], outline=D0ACAC, width=BORDER)
        y += im.height + GAP
    canvas.save(OUT / f"{stem}_composite.png")
print(f"DONE: {len(groups)} composites -> {OUT}", flush=True)