"""Clean an icon sprite sheet: the paper background becomes transparent, and each icon is
centred by its own bounds at one shared size in a square cell.

usage: python3 -I clean_sprites.py <in.png> <out.png> [--cols 3] [--rows 1] [--cell 512] [--fill 0.78]
"""
import argparse
import numpy as np
from PIL import Image, ImageFilter

ap = argparse.ArgumentParser()
ap.add_argument("src"); ap.add_argument("out")
ap.add_argument("--cols", type=int, default=3); ap.add_argument("--rows", type=int, default=1)
ap.add_argument("--cell", type=int, default=512); ap.add_argument("--fill", type=float, default=0.78)
o = ap.parse_args()

img = np.asarray(Image.open(o.src).convert("RGB")).astype(np.float32)
h, w, _ = img.shape
# paper colour: the median of a thin border all round the sheet
border = np.concatenate([img[:6].reshape(-1, 3), img[-6:].reshape(-1, 3), img[:, :6].reshape(-1, 3), img[:, -6:].reshape(-1, 3)])
bg = np.median(border, axis=0)

# alpha from distance to paper; the ramp starts above the paper's own noise, and a 3px median
# drops the isolated specks a generated "flat" paper still has
dist = np.sqrt(((img - bg) ** 2).sum(axis=2))
a = np.clip((dist - 16) / 48, 0, 1)
a = np.asarray(Image.fromarray((a * 255).astype(np.uint8)).filter(ImageFilter.MedianFilter(3))).astype(np.float32) / 255
# recover the ink colour under partial alpha (p = c*a + bg*(1-a)) only where alpha is firm;
# dividing by a small alpha turns noise into bright fringes, so faint edges keep their own colour
safe = np.maximum(a, 1e-3)[..., None]
col = np.where(a[..., None] > 0.5, np.clip((img - bg * (1 - a[..., None])) / safe, 0, 255), img)
rgba = np.dstack([col, a * 255]).astype(np.uint8)

def spans(counts, n, length):
    """Cut one axis into n cells. Each cut starts at the even split and slides, within 20% of a
    cell, to the emptiest line nearby, so an icon that crosses the even line stays whole."""
    step = length / n
    reach = int(step * 0.2)
    cuts = [0]
    for k in range(1, n):
        mid = round(k * step)
        lo, hi = max(1, mid - reach), min(length - 1, mid + reach)
        window = counts[lo:hi]
        best = window.min()
        options = [lo + i for i, v in enumerate(window) if v == best]
        cuts.append(min(options, key=lambda x: abs(x - mid)))
    cuts.append(length)
    return [(cuts[k], cuts[k + 1]) for k in range(n)]

# where the drawing is, for cutting and centring: a 5px median ignores lone specks, so a stray
# dot can't pull an icon off centre (the picture itself keeps every pixel)
drawn = np.asarray(Image.fromarray(((a > 0.1) * 255).astype(np.uint8)).filter(ImageFilter.MedianFilter(5))) > 0
row_spans = spans(drawn.sum(axis=1), o.rows, h)

sheet = Image.new("RGBA", (o.cell * o.cols, o.cell * o.rows), (0, 0, 0, 0))
box = o.cell * o.fill
for r in range(o.rows):
    y0, y1 = row_spans[r]
    col_spans = spans(drawn[y0:y1].sum(axis=0), o.cols, w)
    for c in range(o.cols):
        x0, x1 = col_spans[c]
        part = rgba[y0:y1, x0:x1]
        ys, xs = np.nonzero(drawn[y0:y1, x0:x1])
        n = r * o.cols + c + 1
        if not len(xs):
            raise SystemExit(f"cell {n}: nothing drawn")
        # an icon touching its cell edge was probably cut by a grid that doesn't match the picture
        if xs.min() == 0 or ys.min() == 0 or xs.max() == part.shape[1] - 1 or ys.max() == part.shape[0] - 1:
            print(f"cell {n}: icon touches the cell edge, check the grid")
        icon = Image.fromarray(part[ys.min():ys.max() + 1, xs.min():xs.max() + 1])
        k = box / max(icon.width, icon.height)
        icon = icon.resize((max(1, round(icon.width * k)), max(1, round(icon.height * k))), Image.Resampling.LANCZOS)
        sheet.alpha_composite(icon, (c * o.cell + (o.cell - icon.width) // 2, r * o.cell + (o.cell - icon.height) // 2))
sheet.save(o.out)
print("paper", bg.round().tolist(), "->", o.out, f"({o.cols}x{o.rows})")