#!/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.

`--ab` (R39, the two magpies) lays the distance strip out as a GRID — one line per candidate, the
same view and the same pixel size stacked vertically — because a single 27-cell line is a list and
what a ruling needs is a column you can look down. It then repeats that grid under linear MOTION
BLUR, since the bird is seen for ~1.2 s crossing peripheral vision and a marking that only reads on
a still is not a marking the player ever sees.
"""
import sys
import textwrap
import numpy as np
from PIL import Image, ImageDraw

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


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)
    a = im.getchannel("A")
    bb = a.getbbox()                      # the SUBJECT's extent — motion blur is scaled to the bird
    out.info["subject"] = bb or (0, 0, im.width, im.height)
    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 motion_blur(im, frac):
    """Linear horizontal smear whose length is `frac` × the SUBJECT's on-screen width.

    A true average of N sub-frame positions, not a gaussian: that is what a shutter integrates and
    it is what decides whether a 5 cm white wing bar survives being dragged across its own body.
    """
    bb = im.info.get("subject", (0, 0, im.width, im.height))
    span = max(1, int(round(frac * (bb[2] - bb[0]))))
    if span < 2:
        return im
    pad = span // 2 + 2
    a = np.asarray(Image.new("RGB", (im.width + 2 * pad, im.height), SKY), dtype=np.float32).copy()
    a[:, pad:pad + im.width] = np.asarray(im, dtype=np.float32)
    n = min(96, span)
    acc = np.zeros_like(a)
    for i in range(n):
        off = int(round((i / (n - 1) - 0.5) * span))
        acc += np.roll(a, off, axis=1)
    acc /= n
    out = Image.fromarray(acc[:, pad:pad + im.width].astype("uint8"))
    out.info["subject"] = bb
    return out


def ab_sheet(out_path, title, rows, notes):
    """One line per candidate, stacked, still and then moving. Built for a ruling, not a gallery."""
    w, h = next(iter(rows[0][1].values())).size
    ncol = max(len(v) for _, v in rows)
    cols = [(n, s) for n in STRIP for s in SIZES]
    blocks = [("AT THE SIZE IT IS ACTUALLY SEEN — still (rendered px, magnified NEAREST)", 0.0),
              ("MOVING — smear = 25% of the bird's width (the approach: ~0.11 m/frame at 60 fps)", 0.25),
              ("MOVING — smear = 100% of the bird's width (the pass: swoop covers 34 m in 1.25 s "
               "⇒ 0.45 m/frame, ~1 body length)", 1.0)]
    W = max(PAD + ncol * (w + PAD), GUTTER + len(cols) * (CELL + PAD) + PAD)
    H = (24 + len(rows) * (h + LBL + PAD)
         + len(blocks) * (LBL + 14 + len(rows) * (CELL + LBL + 4)) + PAD
         + LBL * (sum(max(1, len(l) // 300 + 1) for l in notes) * 3 + 2))
    sheet = Image.new("RGB", (W, H), (24, 24, 26))
    d = ImageDraw.Draw(sheet)
    d.text((PAD, 5), title, fill=(240, 240, 240))
    y = 24
    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
    for head, frac in blocks:
        d.text((PAD, y), head, fill=(250, 210, 120))
        y += LBL
        for i, (n, s) in enumerate(cols):
            d.text((GUTTER + i * (CELL + PAD) + 3, y), f"{n} {s}px", fill=(175, 175, 175))
        y += 14
        for label, views in rows:
            d.text((PAD, y + CELL // 2 - 6), label[:22], fill=(225, 225, 225))
            for i, (n, s) in enumerate(cols):
                if n not in views:
                    continue
                src = motion_blur(views[n], frac) if frac else views[n]
                sheet.paste(src.resize((s, s), Image.LANCZOS).resize((CELL, CELL), Image.NEAREST),
                            (GUTTER + i * (CELL + PAD), y))
            y += CELL + LBL + 4
    y += 4
    for line in notes:
        for part in textwrap.wrap(line, width=max(60, (W - 2 * PAD) // 6)) or [""]:
            d.text((PAD, y), part, fill=(190, 190, 190))
            y += LBL
    sheet = sheet.crop((0, 0, W, min(sheet.height, y + PAD)))   # the height estimate over-reserves
    sheet.save(out_path)
    print(f"ab sheet -> {out_path}  {sheet.size}")


def main():
    out_path, title = sys.argv[1], sys.argv[2]
    rows = []
    args = sys.argv[3:]
    extra = None
    notes = []
    if "--note" in args:                       # --note "line" ["line" ...] : footer, always last
        i = args.index("--note")
        notes = args[i + 1:]
        args = args[:i]
    if "--ab" in args:
        args.remove("--ab")
        for i in range(0, len(args), 2):
            rows.append((args[i], load(args[i + 1])))
        return ab_sheet(out_path, title, rows, notes)
    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()