# _stride.py — draw the missing HALF of a stride pair. One frame in, its leg-swapped twin out.
#
#   uv run --with pillow --with numpy --with requests python assets/_stride.py z_rw1 z_rw2
#   uv run --with pillow --with numpy --with requests python assets/_stride.py z_rw1 --check
#
# WHY. A walk cycle in this game is ONE DRAWING WITH THE LEGS MOVED, and every `…b` frame is that
# second half. The club pro's wound rows never got one: `z_rw1b` has the same legs as `z_rw1` (15.1%
# silhouette difference in the bottom third, where the healthy pair scores 51.1%), so a wounded
# opponent has slid through both damage stages since the day it shipped — and because the family
# pipeline draws every new body off the club pro, all seven families inherited it. Oliver, from the
# sprite bench: "arm off — feet are not moving. headless — feet are not moving."
#
# THE PROMPT SHOWS THE CHANGE RATHER THAN DESCRIBING IT. Every previous attempt at this asked in
# words ("identical above the hips, legs clearly swapped") and got two near-identical drawings back.
# IMAGE 2 here is the HEALTHY pair side by side — the same character performing exactly the edit being
# asked for. That is this project's own rule: a picture beats a paragraph.
#
# THE GATE IS THE POINT. A roll is installed only if the legs moved, the body above the hips did not,
# and the racket kept its bed. All three are measured and printed; a failure leaves the sprite alone.
import base64, os, sys
import numpy as np
import requests
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _fine import key, block_size, frames, split_piece
from _famref import legs, stride_diff

HERE = os.path.dirname(os.path.abspath(__file__))
SCRATCH = os.path.join(HERE, '_refcheck')
MODEL = os.environ.get('SPLAT_MODEL', 'google/gemini-3-pro-image')
# THE THRESHOLDS ARE CALIBRATED ON ART THAT ALREADY SHIPS, not chosen. Every accepted stride pair in
# the repo, measured by the two functions below:
#
#   pair              legs    upper        pair              legs    upper
#   z_r2 / z_r3       51.1%    4.4%        z_e1 / z_e2       72.1%   24.0%
#   z_a1 / z_a2       39.2%   13.2%        z_g1 / z_g2       33.6%    8.9%
#   z_b1 / z_b2       58.3%    8.0%        z_br2 / z_br3     40.9%    2.2%
#   z_d1 / z_d2       53.1%    9.9%        z_d3 / z_d4       34.6%    5.7%
#
# ...against the frozen pairs this file exists to replace: z_rw1/z_rw1b 15.1%, z_rw2/z_rw2b 14.7%.
# So a real stride starts at 33.6 and a frozen one tops out at 15.1 — 30 sits in the gap with room on
# both sides. Upper allows 25 because the ball boy's accepted pair spends 24 of it: a step moves the
# torso and the racket arm a little, and on a one-armed or headless body there is less silhouette for
# the same absolute shift to hide in. My first guesses (35 / 14) rejected both of those shipping pairs.
# The two-frame BOSS walk (16.3%) is deliberately outside this: it is a slow pace, not a stride pair.
LEGS_MIN = 30.0        # % the legs must differ
UPPER_MAX = 25.0       # % the body above the hips may differ — it is meant to be the same drawing
BED_MIN = 0.60         # the racket bed must keep at least this fraction of the base frame's bed area


def upper(im):
    """Silhouette of the TOP 60%, normalised like legs() — the half that must NOT change."""
    a = np.array(im.convert('RGBA'))
    op = a[..., 3] > 40
    if not op.any(): return np.zeros((32, 64), bool)
    h, w = op.shape
    y1 = int(h * 0.60)
    ys, xs = np.nonzero(op[:y1])
    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((np.arange(H) * (y1 / H)).round().astype(int), 0, h - 1)
    return op[np.ix_(gy, gx)]


def pair_ref(a, b, out, scale=4):
    """IMAGE 2: two frames side by side on magenta — the edit, demonstrated."""
    ims = [Image.open(os.path.join(HERE, n + '.png')).convert('RGBA') for n in (a, b)]
    ims = [i.resize((i.width * scale, i.height * scale), Image.NEAREST) for i in ims]
    gap = round(ims[0].height * 0.10)
    pad = gap
    W = sum(i.width for i in ims) + gap + 2 * pad
    H = max(i.height for i in ims) + 2 * pad
    s = Image.new('RGB', (W, H), (255, 0, 255))
    x = pad
    for i in ims:
        s.paste(i, (x, H - pad - i.height), i); x += i.width + gap
    s.save(out)
    return out


def one_ref(n, out, scale=4):
    im = Image.open(os.path.join(HERE, n + '.png')).convert('RGBA')
    big = im.resize((im.width * scale, im.height * scale), Image.NEAREST)
    pad = round(big.height * 0.10)
    s = Image.new('RGB', (big.width + 2 * pad, big.height + 2 * pad), (255, 0, 255))
    s.paste(big, (pad, pad), big)
    s.save(out)
    return out


PROMPT = (
 'IMAGE 1 is one figure from a zombie tennis game — a player or a walking corpse, holding a '
 'racket, drawn as hard-edged pixel art on a flat magenta background.\n\n'
 'Redraw IMAGE 1 with ONE change: THE LEGS ARE IN THE OPPOSITE PHASE OF THE SAME STEP. '
 'The leg that is forward becomes the back leg and the leg that is back swings forward, so the two '
 'drawings together read as one body taking a step. The feet must clearly be in different places '
 'than in IMAGE 1.\n\n'
 'IMAGE 2 shows exactly the change being asked for: it is the same character, undamaged, drawn '
 'twice — one drawing, the legs in the two phases of one step. Copy THAT relationship. The '
 'difference between your output and IMAGE 1 must be the difference between the two figures in '
 'IMAGE 2, and nothing else.\n\n'
 'EVERYTHING FROM THE HIPS UP IS UNCHANGED AND IDENTICAL TO IMAGE 1, pixel for pixel where it can '
 'be: the same head or the same headless neck stump, the same shoulders, the same arms, the same '
 'racket in the same hand at the same angle, the same clothing with the same tears and the same '
 'stains, the same wounds with the same blood, the same colours, the same facing, the same height, '
 'the same build. Do NOT redraw the character. Do NOT heal any injury. Do NOT add a head or an arm '
 'that IMAGE 1 does not have. Only the legs and the feet move.\n\n'
 'Match IMAGE 1\'s drawing technique exactly: the same square pixel grid at the same pixel size, the '
 'same hard black outline, the same flat blocks of colour, the same stepped aliased edges, the same '
 'dithering. No soft brushwork, no airbrushing, no blur, no gradients, no anti-aliasing.\n\n'
 'Output ONE figure only, standing on the same ground line, at the same size as IMAGE 1, on a '
 'background of ONE flat solid magenta #FF00FF edge to edge. The inside of the racket head is that '
 'same EXACT #FF00FF. No text, no shadows, no ground lines, no second figure.')


def api_key():
    if os.environ.get('OPENROUTER_API_KEY'): return os.environ['OPENROUTER_API_KEY']
    for line in open(os.path.join(HERE, '..', '.env')):
        if line.startswith('OPENROUTER_API_KEY='):
            return line.split('=', 1)[1].strip().strip('"').strip("'")
    sys.exit('no OPENROUTER_API_KEY')


def u(p):
    return 'data:image/png;base64,' + base64.b64encode(open(p, 'rb').read()).decode()


def pale_blob(im):
    """Largest flat pale blob as a % of the body — the racket bed. A bed drawn as STRINGS breaks into
    fragments and scores a third of what a flat one does, which is the whole use of the number."""
    a = np.array(im.convert('RGBA')).astype(int) if isinstance(im, Image.Image) else im.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; st = [(y, x)]; lab[y, x] = n
        while st:
            cy, cx = st.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; st.append((ny, nx))
    if not n: return 0.0, None, None
    sizes = [(int((lab == i).sum()), i) for i in range(1, n + 1)]
    big, idx = max(sizes)
    m = lab == idx
    colour = a[m][:, :3].mean(0) if m.any() else None
    return round(big / max(1, int((A > 40).sum())) * 100, 2), m, colour


def fill_bed(crop, ref):
    """Repaint the racket's enclosed background in the BASE frame's bed colour.

    `_fine.key` keeps background trapped inside a figure as ink rather than punching it out, so the
    hole is here to be filled; the colour comes off the base's own bed so the pair cannot disagree.
    """
    _, _, colour = pale_blob(ref)
    if colour is None: return crop
    a = crop.copy()
    op = a[..., 3] > 40
    h, w = op.shape
    seen = np.zeros_like(op)
    stack = [(y, x) for x in range(w) for y in (0, h - 1) if not op[y, x]]
    stack += [(y, x) for y in range(h) for x in (0, w - 1) if not op[y, x]]
    for y, x in stack: seen[y, x] = True
    while stack:
        cy, cx = stack.pop()
        for dy, dx in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            ny, nx = cy + dy, cx + dx
            if 0 <= ny < h and 0 <= nx < w and not op[ny, nx] and not seen[ny, nx]:
                seen[ny, nx] = True; stack.append((ny, nx))
    inside = ~op & ~seen
    if inside.any():
        a[inside, :3] = colour.astype(a.dtype)
        a[inside, 3] = 255
    return a


def cut(sheet, like):
    """Slice the roll down to ONE sprite at the SAME height as the frame it twins, on its palette.

    Same height is not a nicety: the game draws every frame at its own PNG height, so a twin that
    came back 20 px taller would make the body rise once per step — the bounce this project has
    already fixed twice.
    """
    src = Image.open(sheet).convert('RGB')
    a = key(src)
    # THE RACKET BED IS FILLED HERE, NOT KEYED OUT — the opposite of what _zsheet.py does, and for the
    # opposite reason. _zsheet cuts a whole family, where every frame's bed keys out together and the
    # racket is an open hoop throughout. A stride twin is half of a PAIR whose other half is already
    # on disk with a painted bed, so keying this one out makes the racket face flicker empty/filled
    # once per step — measured: z_rw1 has 0 px of enclosed transparency and the first z_rw1b roll had
    # 959. `_fine.key` already keeps enclosed background as ink; it is repainted in the base frame's
    # own bed colour so the pair matches.
    b = block_size(a)
    spans = frames(a, [], 1)
    if len(spans) != 1:
        return None, f'{len(spans)} figures in the roll, expected 1'
    x0, x1 = spans[0]
    sub = a[:, x0:x1]
    ys = np.nonzero((sub[..., 3] > 0).any(1))[0]
    xs = np.nonzero((sub[..., 3] > 0).any(0))[0]
    crop = sub[ys.min():ys.max() + 1, xs.min():xs.max() + 1]
    if split_piece(crop):
        return None, 'the roll holds two stacked bodies'
    ref = Image.open(os.path.join(HERE, like + '.png')).convert('RGBA')
    crop = fill_bed(crop, ref)
    snap = Image.fromarray(crop, 'RGBA').resize(
        (max(1, round(crop.shape[1] / b)), max(1, round(crop.shape[0] / b))), Image.BOX)
    k = ref.height / snap.height
    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
    px = np.array(ref)[np.array(ref)[..., 3] > 200][:, :3].astype(int)
    pal = np.unique((px // 8) * 8, axis=0)
    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)
    return Image.fromarray(arr.astype('uint8'), 'RGBA'), None


def gate(new, base):
    """Legs moved, body did not, racket still has a bed. All three numbers, always printed."""
    ref = Image.open(os.path.join(HERE, base + '.png'))
    dl = stride_diff(legs(new), legs(ref))
    du = stride_diff(upper(new), upper(ref))
    bn, br = pale_blob(new)[0], pale_blob(ref)[0]
    # the bed is drawn, not keyed out, so it must survive at close to the base's area. A string mesh
    # fragments it — the first z_rw2b roll scored 2.25 against its base's 7.22.
    bed = bn / br if br else 1.0
    ok = dl >= LEGS_MIN and du <= UPPER_MAX and bed >= BED_MIN
    print(f'    legs {dl:5.1f}% (>={LEGS_MIN})   upper {du:5.1f}% (<={UPPER_MAX})   '
          f'bed {bn:5.2f}% vs {br:5.2f}% = {bed:.2f}x (>={BED_MIN})   {"PASS" if ok else "REJECT"}')
    return ok


def attempt(base, example, i):
    """One roll. EVERY candidate is kept under its own name — a re-roll that comes back worse must
    never be able to overwrite one that passed, which is exactly how the first good z_rw1b was lost."""
    twin = base + 'b'
    i1 = one_ref(base, os.path.join(SCRATCH, f'stride_in_{base}.png'))
    i2 = pair_ref(example[0], example[1], os.path.join(SCRATCH, f'stride_ex_{example[0]}.png'))
    r = requests.post('https://openrouter.ai/api/v1/images',
        headers={'Authorization': f'Bearer {api_key()}', 'Content-Type': 'application/json'},
        json={'model': MODEL, 'prompt': PROMPT, 'resolution': '2K', 'aspect_ratio': '1:1',
              'input_references': [{'type': 'image_url', 'image_url': {'url': u(i1)}},
                                   {'type': 'image_url', 'image_url': {'url': u(i2)}}]},
        timeout=1800)
    if r.status_code != 200:
        print(f'  [{i}] HTTP {r.status_code} {r.text[:160]}'); return None
    d = r.json().get('data') or []
    if not d:
        print(f'  [{i}] no image returned'); return None
    b64 = d[0].get('b64_json') or d[0].get('image_url', {}).get('url', '').split(',')[-1]
    raw = os.path.join(SCRATCH, f'stride_raw_{twin}_{i}.png')
    open(raw, 'wb').write(base64.b64decode(b64))
    im, err = cut(raw, base)
    if err:
        print(f'  [{i}] {err} — discarded'); return None
    cand = os.path.join(SCRATCH, f'stride_cand_{twin}_{i}.png')
    im.save(cand)
    print(f'  [{i}] cut {im.width}x{im.height} -> {cand}')
    return im if gate(im, base) else None


def run(base, example=None, tries=3, install=True):
    fam = base[2:base.index('w')] if 'w' in base[2:] else base[2:]
    example = example or (f'z_{fam}2', f'z_{fam}3')
    twin = base + 'b'
    print(f'{base} -> {twin}   stride example: {example[0]} / {example[1]}')
    for i in range(1, tries + 1):
        im = attempt(base, example, i)
        if im:
            if install:
                im.save(os.path.join(HERE, twin + '.png'))
                print(f'  installed {twin}.png')
            return True
    print(f'  {tries} rolls, none passed — {twin}.png left alone. Candidates are in _refcheck/.')
    return False


if __name__ == '__main__':
    # a flag's VALUE is not a frame name — skip the token after every option that takes one
    argv, args, skip = sys.argv[1:], [], False
    for i, a in enumerate(argv):
        if skip: skip = False; continue
        if a in ('--example', '--tries'): skip = True; continue
        if not a.startswith('-'): args.append(a)
    if not args: sys.exit('usage: _stride.py <frame>... [--example z_rw1] [--tries N] [--check]')
    # --example NAME uses NAME / NAMEb as the demonstration pair instead of the healthy shuffle. A
    # wound stage is a closer picture of the edit than an undamaged body is — same chaining logic the
    # wound ladder itself uses (see _master.py).
    ex = None
    if '--example' in sys.argv:
        e = sys.argv[sys.argv.index('--example') + 1]
        ex = (e, e + 'b')
    tries = int(sys.argv[sys.argv.index('--tries') + 1]) if '--tries' in sys.argv else 3
    if '--check' in sys.argv:
        for base in args:
            print(base)
            gate(Image.open(os.path.join(HERE, base + 'b.png')), base)
    else:
        for base in args:
            run(base, ex, tries, install='--no-install' not in sys.argv)