#!/usr/bin/env python3
"""Compose render_views.py output into one labelled contact sheet + a distance strip.

    PY=~/Documents/MODELBEAST/venvs/mflux/bin/python
    $PY pipeline/view_sheet.py OUT.png "title"  "row label" a.png.views.txt  ["row 2" b...txt] ...

One row of views per input (each composited over sky — bird-against-sky is the contrast that
matters), then a strip that re-samples the swoop / flank / rear views to 64 / 48 / 32 px and
magnifies them NEAREST: a 0.36 m bird at 3-10 m on a 1080-high screen is 50-100 px, so 64 px is
the honest test of whether a marking reads and 32 px is the pessimistic one.
"""
import sys
from PIL import Image, ImageDraw

SKY = (158, 184, 219)
PAD, LBL = 10, 16
STRIP = ("swoop", "flank", "rear34")
SIZES = (64, 48, 32)
CELL = 112


def over_sky(p, bg=SKY):
    im = Image.open(p).convert("RGBA")
    out = Image.new("RGB", im.size, bg)
    out.paste(im, (0, 0), im)
    return out


def load(path):
    views = {}
    for line in open(path):
        k, v = line.rstrip("\n").split("\t")
        if k != "LABEL":
            views[k] = over_sky(v)
    return views


def main():
    out_path, title = sys.argv[1], sys.argv[2]
    rows = []
    args = sys.argv[3:]
    extra = None
    if "--extra" in args:                      # --extra IMG.png "caption": pasted under the sheet
        i = args.index("--extra")
        extra = (args[i + 1], args[i + 2])
        args = args[:i]
    for i in range(0, len(args), 2):
        rows.append((args[i], load(args[i + 1])))
    w, h = next(iter(rows[0][1].values())).size
    ncol = max(len(v) for _, v in rows)
    n_strip = sum(len([n for n in STRIP if n in v]) * len(SIZES) for _, v in rows)
    ex = Image.open(extra[0]).convert("RGB") if extra else None
    W = max(PAD + ncol * (w + PAD), PAD + n_strip * (CELL + PAD), (ex.width + 2 * PAD) if ex else 0)
    H = 20 + len(rows) * (h + LBL + PAD) + LBL + CELL + 18 + PAD
    if ex:
        H += LBL + ex.height + PAD
    sheet = Image.new("RGB", (W, H), (24, 24, 26))
    d = ImageDraw.Draw(sheet)
    d.text((PAD, 4), title, fill=(240, 240, 240))
    y = 22
    for label, views in rows:
        d.text((PAD, y), label, fill=(250, 210, 120))
        y += LBL
        for i, (n, im) in enumerate(views.items()):
            x = PAD + i * (w + PAD)
            sheet.paste(im, (x, y))
            d.text((x + 4, y + 2), n, fill=(40, 40, 40))
        y += h + PAD
    d.text((PAD, y), "AT THE SIZE IT IS ACTUALLY SEEN (rendered px, magnified NEAREST):",
           fill=(250, 210, 120))
    y += LBL
    i = 0
    for label, views in rows:
        for n in STRIP:
            if n not in views:
                continue
            for s in SIZES:
                x = PAD + i * (CELL + PAD)
                sheet.paste(views[n].resize((s, s), Image.LANCZOS).resize((CELL, CELL), Image.NEAREST),
                            (x, y))
                d.text((x + 3, y + CELL + 2), f"{label[:9]} {n} {s}px", fill=(175, 175, 175))
                i += 1
    if ex:
        y += CELL + 18
        d.text((PAD, y), extra[1], fill=(250, 210, 120))
        sheet.paste(ex, (PAD, y + LBL))
    sheet.save(out_path)
    print(f"sheet -> {out_path}  {sheet.size}")


main()