"""Contact sheet of screenshots: montage.py out.png cols label1=dir1 [label2=dir2 ...] --ids v01 v03 ...
With several directories, each row is one task and each column one source (side-by-side comparison)."""
import argparse
import os

from PIL import Image, ImageDraw

ap = argparse.ArgumentParser()
ap.add_argument("out")
ap.add_argument("sources", nargs="+", help="label=dir with <id>_0.png files")
ap.add_argument("--ids", nargs="+", required=True)
ap.add_argument("--w", type=int, default=384)
a = ap.parse_args()
srcs = [s.split("=", 1) for s in a.sources]
w, h, pad = a.w, a.w * 3 // 4, 22
sheet = Image.new("RGB", (len(srcs) * (w + 6) + 6, len(a.ids) * (h + pad + 6) + 6), "#202124")
d = ImageDraw.Draw(sheet)
for r, tid in enumerate(a.ids):
    for c, (label, folder) in enumerate(srcs):
        x, y = 6 + c * (w + 6), 6 + r * (h + pad + 6)
        p = os.path.join(folder, f"{tid}_0.png")
        d.text((x + 4, y + 4), f"{tid}  {label}", fill="#e8eaed")
        if os.path.exists(p):
            sheet.paste(Image.open(p).convert("RGB").resize((w, h)), (x, y + pad))
        else:
            d.rectangle([x, y + pad, x + w, y + pad + h], outline="#5f6368")
            d.text((x + w // 2 - 30, y + pad + h // 2), "no output", fill="#9aa0a6")
sheet.save(a.out)
print(a.out, sheet.size)