# _famref.py — compose THE POSE REFERENCE, `assets/_pose_ref_family.png`, from the installed z_r*
# frames. Locally. No API call.
#
#   uv run --with pillow --with numpy python assets/_famref.py            # write it, then gate it
#   uv run --with pillow --with numpy python assets/_famref.py --check    # gate the file on disk only
#
# WHY THIS EXISTS. The reference every zombie tennis family is drawn from used to be a GENERATED
# RE-DRAW of the racket family: nine separate 4K calls, one per row band, stacked locally
# (`gen4k.py`, session "Zombie sprite regeneration"). Both of the faults Oliver found on the sprite
# bench came out of that round trip, and both were then copied faithfully into all six families:
#
#   THE HEAD. The play row's call came back with the shuffle-A figure DECAPITATED. Measured against
#   the sprite it was drawn from, that one figure scaled 2.341 wide and 1.996 tall where all thirty
#   others were proportional to three decimals — 33 sprite px missing off the top, width untouched.
#   The local input to that call (`in_1_play.png`) was exact at 2.000x. The sheet lost the head, not
#   the sprite: `z_r2.png` on disk has always been fine.
#
#   THE SCALE. Nine independent calls returned nine independent sizes — the fall bands at 81% of the
#   standing bands, the w2 serve band at 87%. `_zsheet.py`'s fall1 re-anchor and its walk-frame
#   levelling exist ONLY to undo that, and the levelling is what turned a headless z_sh2 into a
#   headless AND 13% oversized one: it read the missing head as a height error and scaled up.
#
# A reference composed from the sprites cannot have either fault. One scale, applied once, NEAREST;
# every figure is the art that is already on disk. That is the whole fix, and it costs no tokens.
#
# THE GATE IS ON THE INPUT, WHICH IS WHERE IT WAS MISSING. `_zsheet.py` checks what comes BACK from
# the model — band count, figure count, stacked bodies, edges, palette, heights — and nothing ever
# checked the sheet handed TO it. check() below re-reads the file it just wrote with _zsheet's own
# detection and refuses a sheet the cutter would misread or that has lost a figure's head.
import os, sys
import numpy as np
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _zsheet import layout, bands, key                # the same layout and the same row splitter

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, '_pose_ref_family.png')
FAM = 'r'                                             # the club pro is the master. See PROJECT.md.
MAGENTA = (255, 0, 255)
SIZE = (3609, 6416)                                   # 9:16, the aspect the images endpoint snaps to
MARGIN = 0.025                                        # a PROPORTION, never a pixel count — the model
GAP_X, GAP_Y = 60, 60                                 # picks its own resolution (see PROJECT.md)
STRIDE_MIN = 30.0                                     # % leg difference a real stride pair clears — calibrated in _stride.py


def bed_magenta(im):
    """Repaint the racket's string bed magenta, so it cuts to an open hoop downstream.

    `_zsheet.py` deliberately INVERTS _fine.py's enclosed-magenta rule: the bed is meant to key out.
    That only works if the reference SHOWS a magenta bed, because a picture beats a paragraph — the
    prompt asking for one is not enough on its own. The installed z_r* frames have the bed painted
    pale, so it is repainted here. The bed is the largest pale blob in the frame; the runners-up are
    the shoes, at a third of its area or less on every one of the 31 frames.
    """
    a = np.array(im.convert('RGBA')).astype(int)
    R, G, B, A = a[..., 0], a[..., 1], a[..., 2], a[..., 3]
    lo, hi = np.minimum(np.minimum(R, G), B), np.maximum(np.maximum(R, G), B)
    pale = (A > 40) & (lo > 140) & (hi - lo < 50)
    lab = np.zeros(pale.shape, int)
    n = 0
    for y, x in zip(*np.nonzero(pale)):
        if lab[y, x]: continue
        n += 1
        stack = [(y, x)]; lab[y, x] = n
        while stack:
            cy, cx = stack.pop()
            for dy in (-1, 0, 1):
                for dx in (-1, 0, 1):
                    ny, nx = cy + dy, cx + dx
                    if 0 <= ny < pale.shape[0] and 0 <= nx < pale.shape[1] and pale[ny, nx] and not lab[ny, nx]:
                        lab[ny, nx] = n; stack.append((ny, nx))
    if not n: return im
    areas = [(int((lab == i).sum()), i) for i in range(1, n + 1)]
    best = max(areas)[1]
    a[lab == best, :3] = MAGENTA
    return Image.fromarray(a.astype('uint8'), 'RGBA')


def legs(im):
    """Feet-anchored, mass-centred silhouette of the bottom third — the stride comparison surface.

    Normalised so a wider canvas cannot fake a difference: same metres-per-pixel for every frame,
    centred on the silhouette's own horizontal centre of mass, anchored on the ground line.
    """
    a = np.array(im.convert('RGBA'))
    op = a[..., 3] > 40
    if not op.any(): return np.zeros((32, 64), bool)
    h, w = op.shape
    y0 = int(h * 0.66)
    ys, xs = np.nonzero(op[y0:])
    cx = xs.mean() if len(xs) else w / 2
    W, H, mpp = 64, 32, w / 140
    gx = np.clip((cx + (np.arange(W) - W / 2) * mpp).round().astype(int), 0, w - 1)
    gy = np.clip((h - 1 - (H - 1 - np.arange(H)) * ((h - y0) / H)).round().astype(int), 0, h - 1)
    return op[np.ix_(gy, gx)]


def stride_diff(a, b):
    """% of the union silhouette that differs. A real stride pair clears 33; a frozen one sits at 15."""
    u = (a | b).sum()
    return round(float((a ^ b).sum()) / u * 100, 1) if u else 0.0


def compose():
    rows = layout(FAM)
    frames = {n: bed_magenta(Image.open(os.path.join(HERE, n + '.png'))) for row in rows for n in row}
    # ONE scale for all 31 figures, derived from whichever axis binds. This is the entire point of the
    # file: the old sheet had five scales because it was five separate rolls.
    row_w = [sum(frames[n].width for n in row) + GAP_X * (len(row) - 1) for row in rows]
    row_h = [max(frames[n].height for n in row) for row in rows]
    cw, ch = max(row_w), sum(row_h) + GAP_Y * (len(rows) - 1)
    k = min(SIZE[0] * (1 - 2 * MARGIN) / cw, SIZE[1] * (1 - 2 * MARGIN) / ch)

    sheet = Image.new('RGB', SIZE, MAGENTA)
    y = round((SIZE[1] - ch * k) / 2)
    for row, rh in zip(rows, row_h):
        scaled = [(n, round(frames[n].width * k), round(frames[n].height * k)) for n in row]
        x = round((SIZE[0] - (sum(w for _, w, _ in scaled) + round(GAP_X * k) * (len(row) - 1))) / 2)
        ground = y + round(rh * k)                      # ONE ground line per band, bottom-aligned
        for n, w, h in scaled:
            sheet.paste(frames[n].resize((w, h), Image.NEAREST), (x, ground - h), frames[n].resize((w, h), Image.NEAREST))
            x += w + round(GAP_X * k)
        y = ground + round(GAP_Y * k)
    sheet.save(OUT)
    print(f'{OUT}  {sheet.size}  scale {k:.3f}  31 figures, {len(rows)} bands')
    return k


def check():
    """Re-read the written sheet the way _zsheet.py will, and refuse anything it would misread."""
    rows = layout(FAM)
    a = key(Image.open(OUT).convert('RGB'))
    bs = bands(a)
    bad = []
    if len(bs) != len(rows):
        sys.exit(f'GATE FAILED: {len(bs)} row bands, expected {len(rows)}')
    # scale is read off the stance, exactly as _zsheet.py reads it
    K = None
    for (y0, y1), row in zip(bs, rows):
        band = a[y0:y1]
        cols = (band[..., 3] > 0).any(0)
        spans, s = [], None
        for x, v in enumerate(cols):
            if v and s is None: s = x
            elif not v and s is not None:
                if x - s > 20: spans.append((s, x))
                s = None
        if s is not None: spans.append((s, len(cols)))
        if len(spans) != len(row):
            sys.exit(f'GATE FAILED: row {row[0]}: {len(spans)} figures, expected {len(row)}')
        for (x0, x1), n in zip(spans, row):
            sub = band[:, x0:x1]
            ys = np.nonzero((sub[..., 3] > 0).any(1))[0]
            w, h = x1 - x0, int(ys.max() - ys.min() + 1)
            src = Image.open(os.path.join(HERE, n + '.png'))
            if K is None: K = h / src.height
            sw, sh = w / src.width, h / src.height
            # THE HEAD CLIP IS EXACTLY THIS TEST. A figure that lost its head keeps its width and
            # loses height, so width and height stop agreeing — 2.341 against 1.996 is what the old
            # sheet's z_r2 measured. Everything proportional to 2% is a figure that is all there.
            if abs(sw - sh) > 0.02 * K:
                bad.append(f'{n}: w{sw:.3f} h{sh:.3f} NOT PROPORTIONAL — a clipped figure')
            elif abs(sh - K) > 0.02 * K:
                bad.append(f'{n}: scale {sh:.3f} against the sheet\'s {K:.3f} — not one scale')
    if bad:
        sys.exit('GATE FAILED:\n  ' + '\n  '.join(bad))
    print(f'gate: 31 figures, {len(bs)} bands, every one proportional at one scale {K:.3f}')

    # ...and every stride pair, which the old sheet also got wrong. A frozen pair in the REFERENCE is
    # a frozen pair in every family drawn off it — the model copies what it is shown, so this is the
    # one place the check is worth anything. It was a warning while z_rw1b and z_rw2b still had the
    # same legs as the frames they twin; `_stride.py` drew them, so it is a hard gate now.
    L = {n: legs(Image.open(os.path.join(HERE, n + '.png'))) for n in
         ('z_r2', 'z_r3', 'z_rw1', 'z_rw1b', 'z_rw2', 'z_rw2b')}
    frozen = []
    for lbl, p, q in (('healthy 2<->3', 'z_r2', 'z_r3'), ('w1<->w1b', 'z_rw1', 'z_rw1b'),
                      ('w2<->w2b', 'z_rw2', 'z_rw2b')):
        d = stride_diff(L[p], L[q])
        print(f'  stride {lbl:14s} {d:5.1f}%  {"ok" if d >= STRIDE_MIN else "FROZEN — legs barely move"}')
        if d < STRIDE_MIN: frozen.append(f'{p}/{q} at {d:.1f}%')
    if frozen:
        sys.exit('GATE FAILED: frozen stride pair(s): ' + ', '.join(frozen) +
                 f' — need >={STRIDE_MIN}%. Draw the twin with _stride.py before rolling a family.')


if __name__ == '__main__':
    if '--check' not in sys.argv: compose()
    check()