#!/usr/bin/env python3
"""Turn a generated sprite sheet (figures in a grid on a magenta background) into clean
pixel-art sprites: background removed, every figure cut out, sampled down to a fixed pixel
grid and reduced to a palette.

Usage:
  pixelize_sprite.py raw.png out_prefix [--grid ROWSxCOLS | --views N] --size 32
                     [--crop x0,y0,x1,y1] [--colors 16] [--step S] [--feet Y]
                     [--palette sheet.png] [--add-colors R,G,B;R,G,B] [--measure]
                     [--outline-weight W] [--bg auto|magenta|R,G,B] [--split auto|cells]
                     [--row-counts 4,3] [--pick 0,1,2,m5,3] [--no-despeckle]

--grid / --views  layout of the figures in the raw sheet (--views N = one row of N)
--crop            fractions of the image (0..1) to use only a part of it
--size            output cell (size x size pixels)
--step            source pixels per output pixel; default: the tallest or widest figure fills
                  the cell. Give the same body scale to several sheets of one character with
                  --measure (prints every figure's box in source pixels) and an explicit step.
--feet            the output row the feet stand on (default: the bottom row)
--palette         take the colours of an existing sprite sheet instead of choosing new ones
                  (frames of one character keep exactly the same colours)
--add-colors      extra palette colours (e.g. the light iron of a tool the reference lacks)
--outline-weight  weight of near-black pixels when sampling (below 1: thin details survive)
--bg              background key: auto (default; the colours of the sheet's border, only
                  colours close to them are removed, so purple figures on a green sheet
                  survive), magenta (the old fixed magenta/purple test) or R,G,B
--split           auto (default: figures found by the empty gaps between them, falling back to
                  equal cells) or cells (equal cells of the grid)
--no-despeckle    keep single pixels without any neighbour (removed by default)
--row-counts      figures per row when the rows differ (e.g. 4,3 with --grid 2x4)
--pick            build the output row from these figures (indices in reading order, an "m"
                  in front mirrors it), e.g. 0,1,2,m5,3; the default step uses only these
Frame or grid lines between cells (long straight lines in one flat colour) are removed.

Writes out_prefix-<size>.png (cells in the grid's order, row by row; transparent) and
out_prefix-<size>-preview.png (scaled up 8x on a grey background for viewing). Needs Pillow.
"""
import argparse
import math
from collections import Counter

from PIL import Image


def is_background(p):
    """The old fixed test: magenta and its darker/greyer purple variants. Only used with
    --bg magenta (it also removes purple and pink figures)."""
    r, g, b = p[:3]
    magenta = r > g + 50 and b > g + 50
    frame = abs(r - b) < 40 and r > g + 12 and b > g + 12 and r > 55
    return magenta or frame


def _dist2(p, q):
    return (p[0] - q[0]) ** 2 + (p[1] - q[1]) ** 2 + (p[2] - q[2]) ** 2


def _seg_dist2(p, a, b):
    """Squared distance of colour p to the segment a-b (blends of a and b lie on it)."""
    d = [b[i] - a[i] for i in range(3)]
    dd = sum(v * v for v in d) or 1
    t = max(0.0, min(1.0, sum((p[i] - a[i]) * d[i] for i in range(3)) / dd))
    return sum((p[i] - a[i] - t * d[i]) ** 2 for i in range(3))


def background_colours(img):
    """The sheet's background colours, sampled from the image border (a few px inside, so a
    dark frame around the sheet does not count): every bright colour that covers a good part
    of the border. Usually one (flat magenta or green), sometimes two (cells in two shades)."""
    w, h = img.size
    px = img.load()
    samples = []
    for d in (2, 6, 12):
        samples += [px[x, min(d, h - 1)] for x in range(0, w, 2)]
        samples += [px[x, max(h - 1 - d, 0)] for x in range(0, w, 2)]
        samples += [px[min(d, w - 1), y] for y in range(0, h, 2)]
        samples += [px[max(w - 1 - d, 0), y] for y in range(0, h, 2)]
    q = Counter((p[0] // 32, p[1] // 32, p[2] // 32) for p in samples)
    bgs = []
    for k, n in q.most_common():
        if n < 0.08 * len(samples):
            break
        ps = [p for p in samples if (p[0] // 32, p[1] // 32, p[2] // 32) == k]
        m = tuple(sum(p[i] for p in ps) / len(ps) for i in range(3))
        if max(m) >= 120:  # dark border colours are frames or figures, not the background
            bgs.append(m)
    return bgs or [(255.0, 0.0, 255.0)]


BG_COLOURS = []


def bg_tinted(p):
    """True if p has the background's hue (e.g. magenta: red and blue clearly above green):
    blends of the key colour with thin figure parts, which must not become a palette colour."""
    for b in BG_COLOURS:
        on = [p[i] for i in range(3) if b[i] >= 128]
        off = [p[i] for i in range(3) if b[i] < 128]
        if on and off and min(on) > max(off) + 60:
            return True
    return False


def background_mask(img, mode="auto", tol=70):
    """True where the pixel is background. mode: "auto" (colours from the border, only colours
    near them are keyed, so purple figures on green survive), "magenta" (the old fixed test)
    or "R,G,B" (an explicit key colour). Thin frame lines between cells are removed too."""
    w, h = img.size
    px = img.load()
    data = list(img.get_flattened_data())
    if mode == "magenta":
        bg = [is_background(p) for p in data]
    else:
        bgs = background_colours(img) if mode == "auto" else [tuple(float(v) for v in mode.split(","))]
        BG_COLOURS[:] = bgs
        t2 = tol * tol
        cache = {}
        bg = []
        for p in data:
            v = cache.get(p)
            if v is None:
                # the background or a blend of it with the dark outline (anti-aliased edges;
                # only the lighter half, the darker half belongs to the outline)
                v = cache[p] = any(_dist2(p, b) < t2 or (_seg_dist2(p, b, OUTLINE) < 30 * 30
                                                          and _dist2(p, b) < _dist2(p, OUTLINE))
                                   for b in bgs)
            bg.append(v)
    mask = [bg[y * w:(y + 1) * w] for y in range(h)]
    remove_frame_lines(px, mask, w, h, [] if mode == "magenta" else bgs)
    return mask


OUTLINE = (16, 12, 20)


def remove_frame_lines(px, mask, w, h, bgs=()):
    """Frame/grid lines: a column (row) where almost every pixel is one uniform colour that is
    not background, the figures never cover that much of a line in one flat colour. Those
    pixels, similar ones and blends with the background up to 8 px beside them become
    background."""
    def lines(n, m, get, isbg):
        found = []
        for i in range(n):
            fg = [get(i, j) for j in range(m) if not isbg(i, j)]
            if len(fg) < 0.97 * m:
                continue
            med = sorted(fg, key=sum)[len(fg) // 2]
            if sum(1 for p in fg if _dist2(p, med) < 45 * 45) >= 0.97 * m:
                found.append((i, med))
        return found

    cols = lines(w, h, lambda x, y: px[x, y], lambda x, y: mask[y][x])
    rows = lines(h, w, lambda y, x: px[x, y], lambda y, x: mask[y][x])
    def frameish(p, med):
        return _dist2(p, med) < 70 * 70 or any(_seg_dist2(p, b, med) < 35 * 35 for b in bgs)

    for x, med in cols:
        for xx in range(max(0, x - 8), min(w, x + 9)):
            for y in range(h):
                if frameish(px[xx, y], med):
                    mask[y][xx] = True
    for y, med in rows:
        for yy in range(max(0, y - 8), min(h, y + 9)):
            for x in range(w):
                if frameish(px[x, yy], med):
                    mask[yy][x] = True


def split_runs(occ, lo, hi, n, min_gap):
    """Split [lo, hi) into n figure spans from an occupancy profile: runs of occupied
    positions, the closest neighbours merged until n remain. None if fewer than n runs."""
    runs, start = [], None
    for i in range(lo, hi):
        if occ[i] and start is None:
            start = i
        elif not occ[i] and start is not None:
            runs.append([start, i])
            start = None
    if start is not None:
        runs.append([start, hi])
    # glue runs separated by tiny gaps (a detached spear tip, a sparkle)
    merged = []
    for r in runs:
        if merged and r[0] - merged[-1][1] < min_gap:
            merged[-1][1] = r[1]
        else:
            merged.append(r)
    # drop specks far too small to be a figure
    if merged:
        big = max(r[1] - r[0] for r in merged)
        merged = [r for r in merged if r[1] - r[0] >= 0.08 * big]
    while len(merged) > n:
        # merge the pair with the smallest gap
        i = min(range(len(merged) - 1), key=lambda k: merged[k + 1][0] - merged[k][1])
        merged[i] = [merged[i][0], merged[i + 1][1]]
        del merged[i + 1]
    return merged if len(merged) == n else None


def trim_slivers(mask, x0, x1, y0, y1, min_gap):
    """In equal cells a neighbour's spear tip or wing edge can reach over the cell border:
    small runs of columns at the cell's left or right edge, separated from the figure by a gap,
    are cleared."""
    occ = [sum(mask[y][x] for y in range(y0, y1)) > 3 for x in range(x0, x1)]
    runs, start = [], None
    for i, o in enumerate(occ + [False]):
        if o and start is None:
            start = i
        elif not o and start is not None:
            if runs and start - runs[-1][1] < min_gap:
                runs[-1][1] = i
            else:
                runs.append([start, i])
            start = None
    if len(runs) < 2:
        return
    big = max(r[1] - r[0] for r in runs)
    for r in (runs[0], runs[-1]):
        edge = r[0] == 0 or r[1] == len(occ)
        if edge and r[1] - r[0] < 0.25 * big:
            for x in range(x0 + r[0], x0 + r[1]):
                for y in range(y0, y1):
                    mask[y][x] = False


def components(mask, x0, x1, y0, y1, n, f=4):
    """Find n figures in a band by connected parts (on a grid f times coarser): the n biggest
    parts are the figures, smaller parts (a detached halo, a spear tip) join the nearest
    figure, specks are dropped. Returns [(box, test(x, y))] or None."""
    gw, gh = (x1 - x0 + f - 1) // f, (y1 - y0 + f - 1) // f
    occ = [[False] * gw for _ in range(gh)]
    for y in range(y0, y1):
        row = mask[y]
        orow = occ[(y - y0) // f]
        for x in range(x0, x1):
            if row[x]:
                orow[(x - x0) // f] = True
    lab = [[0] * gw for _ in range(gh)]
    parts = []  # [count, gx0, gy0, gx1, gy1]
    for sy in range(gh):
        for sx in range(gw):
            if occ[sy][sx] and not lab[sy][sx]:
                k = len(parts) + 1
                lab[sy][sx] = k
                stack, part = [(sx, sy)], [0, sx, sy, sx + 1, sy + 1]
                while stack:
                    cx, cy = stack.pop()
                    part[0] += 1
                    part[1], part[2] = min(part[1], cx), min(part[2], cy)
                    part[3], part[4] = max(part[3], cx + 1), max(part[4], cy + 1)
                    for dy in (-1, 0, 1):
                        for dx in (-1, 0, 1):
                            nx, ny = cx + dx, cy + dy
                            if 0 <= nx < gw and 0 <= ny < gh and occ[ny][nx] and not lab[ny][nx]:
                                lab[ny][nx] = k
                                stack.append((nx, ny))
                parts.append(part)
    if len(parts) < n:
        return None
    order = sorted(range(len(parts)), key=lambda i: -parts[i][0])
    big = order[:n]
    # the n-th biggest part must be a real figure, not a stray piece of a merged pair
    if parts[big[-1]][0] < 0.25 * parts[big[0]][0]:
        return None
    owner = {i + 1: i for i in big}
    for i in order[n:]:
        if parts[i][0] < 0.01 * parts[big[0]][0]:
            continue
        p = parts[i]
        cx, cy = (p[1] + p[3]) / 2, (p[2] + p[4]) / 2

        def gap(j):
            q = parts[j]
            return max(q[1] - cx, cx - q[3], 0) + max(q[2] - cy, cy - q[4], 0)
        owner[i + 1] = min(big, key=gap)
    figs = []
    for j in sorted(big, key=lambda j: parts[j][1]):
        mine = {k for k, v in owner.items() if v == j}
        ps_ = [parts[k - 1] for k in mine]
        box = (x0 + min(p[1] for p in ps_) * f, y0 + min(p[2] for p in ps_) * f,
               min(x1, x0 + max(p[3] for p in ps_) * f), min(y1, y0 + max(p[4] for p in ps_) * f))

        def test(x, y, mine=mine):
            return x0 <= x < x1 and y0 <= y < y1 and mask[y][x] and lab[(y - y0) // f][(x - x0) // f] in mine
        figs.append((box, test))
    return figs


def valley_split(cnt, n, min_gap):
    """Figures that touch (a pickaxe reaching into the neighbour): split the widest run of
    occupied columns at its thinnest column (in its middle part) until there are n runs."""
    runs, start = [], None
    for i, c in enumerate(cnt + [0]):
        if c > 3 and start is None:
            start = i
        elif c <= 3 and start is not None:
            if runs and start - runs[-1][1] < min_gap:
                runs[-1][1] = i
            else:
                runs.append([start, i])
            start = None
    if not runs or len(runs) > n:
        return None
    while len(runs) < n:
        k = max(range(len(runs)), key=lambda i: runs[i][1] - runs[i][0])
        a, b = runs[k]
        lo, hi = a + (b - a) // 5, b - (b - a) // 5
        cut = min(range(lo, hi), key=lambda x: cnt[x])
        runs[k:k + 1] = [[a, cut], [cut, b]]
    return [tuple(r) for r in runs]


def tight_box(mask, x0, y0, x1, y1, test=None):
    test = test or (lambda x, y: mask[y][x])
    xs = [x for x in range(x0, x1) if sum(test(x, y) for y in range(y0, y1)) > 3]
    ys = [y for y in range(y0, y1) if sum(test(x, y) for x in range(x0, x1)) > 3]
    return (min(xs), min(ys), max(xs) + 1, max(ys) + 1)


FIGURE_TESTS = []  # per box: None (any figure pixel) or a test(x, y) for that figure only


def figure_boxes(img, rows, cols, bg="auto", split="auto", row_counts=None):
    """Background mask and one box per figure. auto split: by the empty columns between the
    figures, else by connected parts (figures whose tools overlap in x), else equal cells."""
    w, h = img.size
    mask = [[not v for v in row] for row in background_mask(img, bg)]
    boxes = []
    FIGURE_TESTS[:] = []
    for r in range(rows):
        y0, y1 = r * h // rows, (r + 1) * h // rows
        n = row_counts[r] if row_counts else cols
        spans = None
        if split == "auto":
            occ = [sum(mask[y][x] for y in range(y0, y1)) > 3 for x in range(w)]
            spans = split_runs(occ, 0, w, n, max(3, w // 200))
            if spans is None:
                figs = components(mask, 0, w, y0, y1, n)
                if figs:
                    for box, test in figs:
                        boxes.append(tight_box(mask, *box, test))
                        FIGURE_TESTS.append(test)
                    continue
                spans = valley_split([sum(mask[y][x] for y in range(y0, y1)) for x in range(w)],
                                     n, max(3, w // 200))
        if spans is None:
            spans = [(c * w // n, (c + 1) * w // n) for c in range(n)]
        for (x0, x1) in spans:
            trim_slivers(mask, x0, x1, y0, y1, max(3, w // 200))
            boxes.append(tight_box(mask, x0, y0, x1, y1))
            FIGURE_TESTS.append(None)
    return mask, boxes


def despeckle(grid):
    """Remove single opaque pixels with no opaque 8-neighbour (stray sparkles, key leftovers)."""
    n = len(grid)
    out = [row[:] for row in grid]
    for y in range(n):
        for x in range(n):
            if grid[y][x] and not any(
                    0 <= y + dy < n and 0 <= x + dx < n and grid[y + dy][x + dx]
                    for dy in (-1, 0, 1) for dx in (-1, 0, 1) if dy or dx):
                out[y][x] = None
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("raw")
    ap.add_argument("out")
    ap.add_argument("--views", type=int, default=5)
    ap.add_argument("--grid", default=None)
    ap.add_argument("--size", type=int, default=24)
    ap.add_argument("--colors", type=int, default=16)
    ap.add_argument("--crop", default=None)
    ap.add_argument("--step", type=float, default=None)
    ap.add_argument("--feet", type=int, default=None)
    ap.add_argument("--palette", default=None)
    ap.add_argument("--add-colors", default="")
    ap.add_argument("--outline-weight", type=float, default=1.0)
    ap.add_argument("--measure", action="store_true")
    ap.add_argument("--bg", default="auto")
    ap.add_argument("--split", default="auto", choices=["auto", "cells"])
    ap.add_argument("--no-despeckle", action="store_true")
    ap.add_argument("--row-counts", default=None)
    ap.add_argument("--pick", default=None)
    a = ap.parse_args()
    rows, cols = (int(v) for v in a.grid.split("x")) if a.grid else (1, a.views)
    feet = a.size - 1 if a.feet is None else a.feet

    img = Image.open(a.raw).convert("RGB")
    if a.crop:
        f = [float(v) for v in a.crop.split(",")]
        w, h = img.size
        img = img.crop((int(f[0] * w), int(f[1] * h), int(f[2] * w), int(f[3] * h)))
    row_counts = [int(v) for v in a.row_counts.split(",")] if a.row_counts else None
    mask, boxes = figure_boxes(img, rows, cols, a.bg, a.split, row_counts)
    if a.measure:
        for i, (x0, y0, x1, y1) in enumerate(boxes):
            print(f"figure {i}: x {x0}-{x1}, y {y0}-{y1}: {x1 - x0} x {y1 - y0} source pixels")
        return
    px = img.load()
    # default: one scale for all figures, the tallest or widest fills the cell
    used_boxes = [boxes[int(t.lstrip("m"))] for t in a.pick.split(",")] if a.pick else boxes
    step = a.step or max(max(b[3] - b[1], b[2] - b[0]) for b in used_boxes) / a.size

    grids = []
    for i, (x0, y0, x1, y1) in enumerate(boxes):
        inside = FIGURE_TESTS[i] or (lambda x, y: mask[y][x])
        gw, gh = round((x1 - x0) / step), round((y1 - y0) / step)
        grid = [[None] * a.size for _ in range(a.size)]
        ox, oy = (a.size - gw) // 2, feet + 1 - gh  # centred, the feet on the feet row
        for gy in range(gh):
            for gx in range(gw):
                # the commonest colour of the block's inner part (edges blur between pixels)
                cx0, cy0 = x0 + gx * step, y0 + gy * step
                votes, bg, fg = Counter(), 0, 0
                for sy in range(4):
                    for sx in range(4):
                        x = int(cx0 + (sx + 0.5) * step / 4)
                        y = int(cy0 + (sy + 0.5) * step / 4)
                        if x >= img.size[0] or y >= img.size[1] or not inside(x, y):
                            bg += 1
                        else:
                            r, g, b = px[x, y]
                            if bg_tinted((r, g, b)):
                                continue  # a blend with the background: no vote at all
                            fg += 1
                            # the outline counts a little less, so thin coloured parts
                            # (a tool's iron, a handle) are not swallowed by their outline
                            votes[(r // 8 * 8, g // 8 * 8, b // 8 * 8)] += a.outline_weight if r + g + b < 90 else 1
                if votes and fg >= bg and 0 <= gy + oy < a.size and 0 <= gx + ox < a.size:
                    grid[gy + oy][gx + ox] = votes.most_common(1)[0][0]
        grids.append(grid if a.no_despeckle else despeckle(grid))

    if a.pick:
        # a new single row from chosen figures; "m" in front mirrors the figure
        picked = []
        for t in a.pick.split(","):
            g = grids[int(t.lstrip("m"))]
            picked.append([row[::-1] for row in g] if t.startswith("m") else g)
        grids, rows, cols = picked, 1, len(picked)
    used = [c for g in grids for row in g for c in row if c]
    counts = Counter(used)
    dist = lambda p, q: sum((x - y) ** 2 for x, y in zip(p, q)) ** 0.5
    if a.palette:
        ref = Image.open(a.palette).convert("RGBA")
        palette = sorted({p[:3] for p in ref.get_flattened_data() if p[3] > 128})
    else:
        # farthest-point picking weighted by how often a colour occurs, so small but distinct
        # details (eyes, the tool's iron) keep their own colour instead of melting into the
        # large areas
        cands = [c for c, n in counts.items() if n >= 3] or list(counts)
        palette = [counts.most_common(1)[0][0]]
        while len(palette) < min(a.colors, len(cands)):
            palette.append(max(cands, key=lambda c: min(dist(c, q) for q in palette) * math.log(1 + counts[c])))
    palette += [tuple(int(v) for v in c.split(",")) for c in a.add_colors.split(";") if c]
    near = {c: min(palette, key=lambda q: dist(c, q)) for c in counts}

    sheet = Image.new("RGBA", (a.size * cols, a.size * rows), (0, 0, 0, 0))
    for i, g in enumerate(grids):
        cx, cy = (i % cols) * a.size, (i // cols) * a.size
        for y, row in enumerate(g):
            for x, c in enumerate(row):
                if c:
                    sheet.putpixel((cx + x, cy + y), near[c] + (255,))
    sheet.save(f"{a.out}-{a.size}.png")
    prev = Image.new("RGBA", sheet.size, (96, 92, 104, 255))
    prev.alpha_composite(sheet)
    prev.resize((sheet.width * 8, sheet.height * 8), Image.NEAREST).save(f"{a.out}-{a.size}-preview.png")
    print(f"{a.out}-{a.size}.png: {rows}x{cols} figures, {a.size}px, step {step:.1f}, {len(set(near.values()))} colours")


if __name__ == "__main__":
    main()