# _fine.py — turn a generated sheet into FULL-RESOLUTION sprites.
#
# _slice.py resamples with NEAREST, which is right when the generator's art-pixel grid is coarse but
# shreds fine detail at a non-integer ratio. This one instead:
#   1. detects the art-pixel BLOCK SIZE the generator actually drew at, and snaps to that grid
#      (one output pixel per art pixel — no half-blocks, no broken outlines),
#   2. common-scales the sheet to the target height with a box filter,
#   3. quantises every pixel back to the sheet's own palette, so the result is flat pixel art
#      rather than a blurred photo of pixel art.
# Usage: _fine.py sheet.png out1.png out2.png ...   [--h 224] [--cut X,Y,...]
import sys
from PIL import Image, ImageDraw
import numpy as np

TARGET_H = 224
KEY_D2 = 60 * 60
MIN_GAP = 12
MIN_BLOB = 60

def key(im, defringe=2, bed=False):
    a = np.array(im.convert('RGBA'))
    bg = a[0, 0, :3].astype(int)
    d = ((a[..., :3].astype(int) - bg) ** 2).sum(2)
    magenta = np.minimum(a[..., 0], a[..., 2]).astype(int) > a[..., 1].astype(int) + 40
    bgish = (d < KEY_D2) | magenta
    # ONLY THE OUTSIDE IS BACKGROUND. This used to punch out every magenta pixel on the sheet
    # wherever it was, and the generator paints a racket's STRINGS in the backdrop colour more often
    # than not — so frame after frame installed a man holding an empty hoop. Background is the region
    # CONNECTED TO THE BORDER (which touches_edge already guarantees no figure reaches), so an
    # enclosed patch of it inside a racket head is not background and does not become a hole.
    # .copy() is load-bearing: an Image.fromarray image is READ-ONLY, and ImageDraw.floodfill's
    # first write raises inside its own try/except and returns having filled nothing — a silent
    # no-op that looks exactly like "the whole sheet is one figure".
    m = Image.fromarray(np.where(bgish, 255, 0).astype('uint8'), 'L').copy()
    ImageDraw.floodfill(m, (0, 0), 128)
    outside = np.array(m) == 128
    if not outside.any(): sys.exit('ERROR: background flood filled nothing — is (0,0) not background?')
    a[outside] = 0
    # ...and what that leaves — magenta INSIDE a body — is repainted as pale grey rather than kept,
    # because bright pink strings are as wrong as no strings. Luminance is preserved, so the drawn
    # crosshatch survives as crosshatch; only the hue goes.
    # UNLESS the sheet MEANS its magenta beds (`bed=True` — the famroll convention, adopted by the
    # opponent refresh and _runroll): then enclosed magenta is the string bed and CUTS TO
    # TRANSPARENCY, exactly as _zsheet already treats it. The grey repaint is what shipped the
    # opponents holding solid white paddles.
    inner = magenta & ~outside
    if bed:
        a[inner] = 0
        inner = np.zeros_like(inner)
    if inner.any():
        lum = a[..., :3].mean(2)
        a[inner, 0] = a[inner, 1] = a[inner, 2] = np.clip(lum[inner] * 0.5 + 115, 0, 255).astype('uint8')
    # De-fringe. The generator does not always give you a pure key colour, and whatever it does give
    # bleeds a halo one pixel into the sprite — which survives as a magenta rim once the sprite is on
    # a green court. Anything on the EDGE that is still leaning towards the background goes too.
    # ...but each pass also eats a 1-pixel line, and a racket frame IS a 1-pixel line: two passes
    # removed the rackets from a whole serve sheet and left men holding a grip. --defringe 1 (or 0)
    # is for sheets whose art is that thin.
    for _ in range(defringe):
        solid = a[..., 3] > 0
        edge = solid & ~(np.roll(solid, 1, 0) & np.roll(solid, -1, 0) & np.roll(solid, 1, 1) & np.roll(solid, -1, 1))
        near = ((a[..., :3].astype(int) - bg) ** 2).sum(2) < KEY_D2 * 9
        a[edge & near] = 0
    return a

def block_size(a):
    ink = a[..., 3] > 0
    ys, xs = np.where(ink)
    sub = a[ys.min():ys.max()+1, xs.min():xs.max()+1]
    h, w = sub.shape[:2]
    best = 1
    for b in range(2, 13):
        if h < b * 6: break
        ok = tot = 0
        for y in range(0, h - b, b):
            for x in range(0, w - b, b):
                blk = sub[y:y+b, x:x+b].reshape(-1, 4)
                # ONLY BLOCKS THAT TOUCH THE ART GET A VOTE. key() zeroes the background, so an
                # empty block is uniform at EVERY b — and on a 9-row family sheet the background is
                # ~85% of the bounding box, so it alone cleared the 0.82 bar and reported b=11 for
                # sheets whose figures were drawn fine. _zsheet then snapped those figures down 11x
                # and back up: that is where the pixelated families came from, not the model.
                if not blk[:, 3].any(): continue
                tot += 1
                if (blk == blk[0]).all(): ok += 1
        if tot and ok / tot > 0.82: best = b
    return best

# WHERE TO CUT WHEN THE FIGURES TOUCH. Segmentation by empty columns needs a clear gap, and the
# generator routinely spaces a row so a belly, an elbow or a racket overlaps its neighbour — a whole
# sheet of otherwise perfect art then reports "found 2 frames, expected 4" and is thrown away. That
# happened five times in one afternoon, and the sheet it lost last was the best roll of the day.
# A row is EVENLY SPACED by construction (we ask for one), so when the count is known the honest
# recovery is: look near each even division and cut at the THINNEST column within that window —
# between two figures the ink thins out even where it does not vanish. Never used to guess a count,
# only to reach a count we already asked for.
def valley_cuts(ink_cols, x0, x1, n):
    span = x1 - x0
    win = max(2, span // (3 * n))
    out = []
    for k in range(1, n):
        c = x0 + span * k // n
        lo, hi = max(x0 + 1, c - win), min(x1 - 1, c + win)
        if hi <= lo: return []
        out.append(lo + int(np.argmin(ink_cols[lo:hi])))
    return out

# DO THE FIGURES SHARE A GROUND LINE? That is the invariant every prompt in this project asks for,
# and it is the one that tells a ROW from a GRID. An empty-band test does not: in the sheet that
# caught this, a 2x2 of a brute falling, the bottom-right figure was tall enough to overlap the top
# row and there was no clean horizontal gap anywhere.
#
# It matters because the valley rescue below will cut ANY sheet into the number of columns it was
# asked for — and on a grid each column is a torso above someone else's boots. Four such frames
# installed silently and looked, in the log, exactly like a success. The rescue turned a loud failure
# into a quiet one, which is the worst trade available in this file; this is the check that stops it.
# IS THIS ONE FIGURE, OR TWO STACKED? Asked of each CUT PIECE, not of the sheet. A body is
# vertically continuous — head to feet there is no height at which the piece is empty — so a
# full-width empty band inside a piece's own extent means the cut caught someone's torso above
# someone else's boots.
#
# Asking it of the whole sheet does not work, and both weaker versions were tried: an empty-band test
# on the sheet fails when a grid's bottom-right figure is tall enough to overlap the top row, and a
# shared-ground-line test fails because after a bad cut EVERY column reaches the bottom row's feet.
# The piece is the only place the question has a clean answer.
def split_piece(crop):
    rows = (crop[..., 3] > 0).any(1)
    ys = np.where(rows)[0]
    if not len(ys): return False
    run = best = 0
    for y in range(ys.min(), ys.max() + 1):
        run = 0 if rows[y] else run + 1
        if run > best: best = run
    return best >= MIN_GAP * 2

def frames(a, cuts, want=0):
    ink = (a[..., 3] > 0)
    cols = ink.any(0)
    out, start, gap = [], None, 0
    for x in range(len(cols)):
        if cols[x]:
            if start is None: start = x
            gap = 0
        elif start is not None:
            gap += 1
            if gap >= MIN_GAP: out.append((start, x - gap + 1)); start = None
    if start is not None: out.append((start, len(cols)))
    for c in sorted(cuts, reverse=True):          # manual splits where two figures touch
        for i, (x0, x1) in enumerate(out):
            if x0 < c < x1: out[i:i+1] = [(x0, c), (c, x1)]; break
    # ...and the automatic version of the same rescue, when no manual cut was given
    if want and len(out) < want and not cuts:
        xs = np.where(cols)[0]
        if len(xs):
            auto = valley_cuts(ink.sum(0), int(xs.min()), int(xs.max()) + 1, want)
            if auto:
                print(f'note: only {len(out)} blob(s) — figures touch; cutting at {auto}')
                edges = [int(xs.min())] + auto + [int(xs.max()) + 1]
                out = list(zip(edges, edges[1:]))
    return out

def main(sheet, outs, target_h, cuts, anchor, defringe=2, bed=False):
    a = key(Image.open(sheet), defringe, bed)
    b = block_size(a)
    crops = []
    spans = frames(a, cuts, len(outs))
    for x0, x1 in spans:
        sub = a[:, x0:x1]
        ys = np.where((sub[..., 3] > 0).any(1))[0]; xs = np.where((sub[..., 3] > 0).any(0))[0]
        crops.append(sub[ys.min():ys.max()+1, xs.min():xs.max()+1])
    if len(crops) != len(outs):
        sys.exit(f'ERROR: found {len(crops)} frames, expected {len(outs)} ({sheet})')
    bad = [i for i, c in enumerate(crops) if split_piece(c)]
    if bad:
        sys.exit(f'ERROR: frame(s) {bad} contain two stacked bodies — the sheet is a GRID, not one '
                 f'row ({sheet})')
    # palette = the colours the generator actually used, on the largest frame
    big = max(crops, key=lambda c: (c[..., 3] > 0).sum())
    px = big[big[..., 3] > 200][:, :3]
    pal = np.unique((px // 8) * 8, axis=0)
    # ONE scale for the sheet, anchored on a chosen frame — the STANDING pose is the 224 reference,
    # so a smash with the arm overhead stays legitimately taller (the game scales off img.height).
    anchor_h = crops[anchor].shape[0] if 0 <= anchor < len(crops) else max(c.shape[0] for c in crops)
    k = target_h / (anchor_h / b)
    for c, out in zip(crops, outs):
        h, w = c.shape[:2]
        snap = Image.fromarray(c, 'RGBA').resize((max(1, round(w / b)), max(1, round(h / b))), Image.BOX)
        fh = max(1, round(snap.height * k)); fw = max(1, round(snap.width * k))
        im = snap.resize((fw, fh), Image.BOX)
        arr = np.array(im).astype(int)
        solid = arr[..., 3] > 110
        flat = arr[..., :3].reshape(-1, 3)
        idx = ((flat[:, None, :] - pal[None, :, :]) ** 2).sum(2).argmin(1)
        arr[..., :3] = pal[idx].reshape(arr.shape[0], arr.shape[1], 3)
        arr[..., 3] = np.where(solid, 255, 0)
        Image.fromarray(arr.astype('uint8'), 'RGBA').save(out)
        print(f'{out}: {fw}x{fh}  (art block {b} in source)')

if __name__ == '__main__':
    args = sys.argv[1:]
    th, cuts, anchor = TARGET_H, [], -1
    if '--h' in args: i = args.index('--h'); th = int(args[i+1]); del args[i:i+2]
    if '--cut' in args: i = args.index('--cut'); cuts = [int(v) for v in args[i+1].split(',')]; del args[i:i+2]
    de = 2
    if '--defringe' in args: i = args.index('--defringe'); de = int(args[i+1]); del args[i:i+2]
    if '--anchor' in args: i = args.index('--anchor'); anchor = int(args[i+1]); del args[i:i+2]
    main(args[0], args[1:], th, cuts, anchor, de)