#!/usr/bin/env python3
"""Tile labeled images into a contact sheet for side-by-side comparison.

Usage: montage.py --out sheet.png --cols 3 [--scale 0.5] path:label path:label ...
"""
import argparse

from PIL import Image, ImageDraw, ImageFont


def load_font(size):
    for p in ("/System/Library/Fonts/Supplemental/Arial.ttf",
              "/System/Library/Fonts/Helvetica.ttc",
              "/Library/Fonts/Arial.ttf"):
        try:
            return ImageFont.truetype(p, size)
        except OSError:
            continue
    return ImageFont.load_default()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--cols", type=int, default=3)
    ap.add_argument("--scale", type=float, default=1.0)
    ap.add_argument("--pad", type=int, default=8)
    ap.add_argument("--label-h", type=int, default=24, dest="label_h")
    ap.add_argument("items", nargs="+", help="each 'path:label'")
    a = ap.parse_args()

    imgs, labels = [], []
    for it in a.items:
        path, _, lab = it.partition(":")
        imgs.append(Image.open(path).convert("RGB"))
        labels.append(lab)

    cw = int(max(i.width for i in imgs) * a.scale)
    ch = int(max(i.height for i in imgs) * a.scale)
    font = load_font(max(12, int(a.label_h * 0.7)))
    cols = a.cols
    rows = (len(imgs) + cols - 1) // cols
    cellw, cellh = cw + a.pad, ch + a.label_h + a.pad
    sheet = Image.new("RGB", (cols * cellw + a.pad, rows * cellh + a.pad), (18, 18, 20))
    d = ImageDraw.Draw(sheet)
    for k, (im, lab) in enumerate(zip(imgs, labels)):
        r, c = divmod(k, cols)
        x, y = a.pad + c * cellw, a.pad + r * cellh
        d.text((x + 2, y + 4), lab, font=font, fill=(232, 232, 238))
        sheet.paste(im.resize((cw, ch), Image.LANCZOS), (x, y + a.label_h))
    sheet.save(a.out)
    print(f"wrote {a.out} ({sheet.width}x{sheet.height})")


if __name__ == "__main__":
    main()