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

Usage:
  pixelize_sprite.py raw.png out_prefix --views 5 --size 24 [--crop x0,y0,x1,y1] [--colors 16]

--crop takes fractions of the image (0..1) to use only a part of it (e.g. one row of a sheet).
Writes out_prefix-<size>.png (the views side by side, size x size each, transparent) and
out_prefix-<size>-preview.png (scaled up 8x on a dark 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, views):
    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 v in range(views):
        x0, x1 = v * w // views, (v + 1) * w // views
        cols = [x for x in range(x0, x1) if sum(mask[y][x] for y in range(h)) > 3]
        rows = [y for y in range(h) if sum(mask[y][x] for x in range(x0, x1)) > 3]
        boxes.append((min(cols), min(rows), max(cols) + 1, max(rows) + 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("--size", type=int, default=24)
    ap.add_argument("--colors", type=int, default=16)
    ap.add_argument("--crop", default=None)
    a = ap.parse_args()

    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, a.views)
    px = img.load()
    # one scale for all views: the tallest or widest figure fills the cell
    step = 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, a.size - gh  # centred, standing on the bottom
        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 = Counter(), 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]
                            votes[(r // 8 * 8, g // 8 * 8, b // 8 * 8)] += 1
                if votes and sum(votes.values()) >= 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)

    # one shared palette for all views: 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
    used = [c for g in grids for row in g for c in row if c]
    counts = Counter(used)
    cands = [c for c, n in counts.items() if n >= 3] or list(counts)
    dist = lambda p, q: sum((x - y) ** 2 for x, y in zip(p, q)) ** 0.5
    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])))
    near = {c: min(palette, key=lambda q: dist(c, q)) for c in counts}

    sheet = Image.new("RGBA", (a.size * a.views, a.size), (0, 0, 0, 0))
    for v, g in enumerate(grids):
        for y, row in enumerate(g):
            for x, c in enumerate(row):
                if c:
                    sheet.putpixel((v * a.size + x, 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: {a.views} views, {a.size}px, {len(set(near.values()))} colours")


if __name__ == "__main__":
    main()