#!/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]

--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)

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):
    r, g, b = p[:3]
    # magenta and its darker/greyer purple variants (sheet background and frame lines)
    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 figure_boxes(img, rows, cols):
    w, h = img.size
    px = img.load()
    mask = [[not is_background(px[x, y]) for x in range(w)] for y in range(h)]
    boxes = []
    for r in range(rows):
        y0, y1 = r * h // rows, (r + 1) * h // rows
        for c in range(cols):
            x0, x1 = c * w // cols, (c + 1) * w // cols
            xs = [x for x in range(x0, x1) if sum(mask[y][x] for y in range(y0, y1)) > 3]
            ys = [y for y in range(y0, y1) if sum(mask[y][x] for x in range(x0, x1)) > 3]
            boxes.append((min(xs), min(ys), max(xs) + 1, max(ys) + 1))
    return mask, boxes


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")
    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)))
    mask, boxes = figure_boxes(img, rows, cols)
    if a.measure:
        for i, (x0, y0, x1, y1) in enumerate(boxes):
            print(f"figure {i} (row {i // cols}, column {i % cols}): {x1 - x0} x {y1 - y0} source pixels")
        return
    px = img.load()
    # default: one scale for all figures, the tallest or widest fills the cell
    step = a.step or max(max(b[3] - b[1], b[2] - b[0]) for b in boxes) / a.size

    grids = []
    for (x0, y0, x1, y1) in boxes:
        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 mask[y][x]:
                            bg += 1
                        else:
                            r, g, b = px[x, y]
                            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)

    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()