# _zsheet.py — ONE generated multi-row sheet -> a whole zombie family of sprites.
#
#   python3 assets/_zsheet.py assets/_refcheck/brute_all3.png br --out assets/_refcheck/sprites
#   python3 assets/_zsheet.py assets/_refcheck/brute_all3.png br --install
#
# WHY THIS EXISTS. _regen.py generates a family as nine separate sheets, and nine separate rolls is
# nine chances for the character to drift — which is exactly what it did, all day. Asking for the
# WHOLE family in one image fixes that by construction: one roll, one character, one light, one
# palette. What was missing was the way back down — a generated sheet is 9 rows of figures and
# _fine.py cuts a single row. This is that step, and nothing else.
#
# TWO THINGS ARE GLOBAL, and both were bugs when they were not:
#   SCALE    _fine.py anchors one scale per sheet, so nine bands would each be scaled to their own
#            tallest figure — the serve strike would come out the same height as the idle stance and
#            every pose would pop, because the game sizes a sprite off its own PNG height. The scale
#            is taken ONCE, from the standing frame of the play row, and every band uses it.
#   PALETTE  taken once from the largest figure in the whole sheet, so a wound row cannot quantise
#            itself onto a slightly different set of colours than the healthy row it follows.
import os, sys
import numpy as np
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _fine import key, block_size, frames, split_piece            # the proven keyer and column cutter

HERE = os.path.dirname(os.path.abspath(__file__))
TARGET_H = 224
MIN_BAND = 40

# The family layout, exactly config.famFrames' order. The wound rows carry no separate stance —
# the shuffle doubles as one — which is why they are 3 and not 4.
def layout(k):
    P = ['bounce', 'toss', 'serve']
    F = ['fall1', 'fall2', 'fall3', 'rise']
    rows = [[f'z_{k}1', f'z_{k}2', f'z_{k}3', f'z_{k}4'],
            [f'z_{k}{s}' for s in P],
            [f'z_{k}{s}' for s in F]]
    for st in (1, 2):
        rows += [[f'z_{k}w{st}', f'z_{k}w{st}b', f'z_{k}w{st}s'],
                 [f'z_{k}w{st}{s}' for s in P],
                 [f'z_{k}w{st}{s}' for s in F]]
    return rows


def bands(a):
    """Row bands = runs of rows that contain ink. Rows are what a grid has and _fine.py refuses;
    here they are the point, so they are found first and each band is handed over as its own row."""
    rows = (a[..., 3] > 0).any(1)
    out, s = [], None
    for y, v in enumerate(rows):
        if v and s is None: s = y
        elif not v and s is not None:
            if y - s >= MIN_BAND: out.append((s, y))
            s = None
    if s is not None and len(rows) - s >= MIN_BAND: out.append((s, len(rows)))
    return out


def cut(sheet, fam, outdir):
    want = layout(fam)
    src = Image.open(sheet).convert('RGB')
    a = key(src)
    # THE ENCLOSED MAGENTA IS THE RACKET, AND HERE IT IS MEANT TO GO. _fine.py repaints magenta that
    # is trapped inside a figure as pale grey, because the generator used to fill string beds with
    # the backdrop colour BY ACCIDENT and keying them left men holding an empty hoop. This sheet is
    # the opposite case: the prompt asks for the string bed to be exactly the background magenta so
    # it cuts to transparency, and it obeyed. So every magenta pixel goes, connected to the border
    # or not. That is safe only because nothing else on the figure is magenta — check a new sheet.
    n = np.array(src).astype(int)
    lo = np.minimum(n[..., 0], n[..., 2])
    a[(lo > 140) & (n[..., 1] < lo - 55), 3] = 0
    b = block_size(a)
    bs = bands(a)
    if len(bs) != len(want):
        sys.exit(f'ERROR: found {len(bs)} row bands, expected {len(want)}')

    crops = []
    for (y0, y1), row in zip(bs, want):
        band = a[y0:y1]
        spans = frames(band, [], len(row))
        if len(spans) != len(row):
            sys.exit(f'ERROR: row {row[0]}: found {len(spans)} figures, expected {len(row)}')
        for x0, x1 in spans:
            sub = band[:, x0:x1]
            ys = np.where((sub[..., 3] > 0).any(1))[0]
            xs = np.where((sub[..., 3] > 0).any(0))[0]
            c = sub[ys.min():ys.max() + 1, xs.min():xs.max() + 1]
            if split_piece(c):
                sys.exit(f'ERROR: a frame in row {row[0]} holds two stacked bodies')
            crops.append(c)

    flat = [n for row in want for n in row]
    # ONE scale, off the play row's standing frame; ONE palette, off the biggest figure anywhere.
    k = TARGET_H / (crops[0].shape[0] / b)
    big = max(crops, key=lambda c: (c[..., 3] > 0).sum())
    px = big[big[..., 3] > 200][:, :3]
    pal = np.unique((px // 8) * 8, axis=0)

    os.makedirs(outdir, exist_ok=True)
    for c, name in zip(crops, flat):
        h, w = c.shape[:2]
        snap = Image.fromarray(c, 'RGBA').resize((max(1, round(w / b)), max(1, round(h / b))), Image.BOX)
        im = snap.resize((max(1, round(snap.width * k)), max(1, round(snap.height * k))), Image.BOX)
        arr = np.array(im).astype(int)
        solid = arr[..., 3] > 110
        f = arr[..., :3].reshape(-1, 3)
        arr[..., :3] = pal[((f[:, None, :] - pal[None, :, :]) ** 2).sum(2).argmin(1)].reshape(arr.shape[0], arr.shape[1], 3)
        arr[..., 3] = np.where(solid, 255, 0)
        Image.fromarray(arr.astype('uint8'), 'RGBA').save(os.path.join(outdir, name + '.png'))
    print(f'{len(flat)} frames, art block {b}, scale {k:.3f}, palette {len(pal)} colours -> {outdir}')
    for (y0, y1), row in zip(bs, want):
        hs = [Image.open(os.path.join(outdir, n + '.png')).height for n in row]
        print(f'  {row[0]:14s} {len(row)} frames  heights {hs}')


if __name__ == '__main__':
    args = [a for a in sys.argv[1:] if not a.startswith('--')]
    out = HERE if '--install' in sys.argv else None
    if '--out' in sys.argv: out = sys.argv[sys.argv.index('--out') + 1]
    cut(args[0], args[1], out or os.path.join(HERE, '_refcheck', 'sprites'))