#!/usr/bin/env python3
"""Cut a generated sheet of icons (a grid on a flat magenta background) into transparent icons
for the high-res atlas: every cell keyed (only nearly pure magenta goes, so violet glows stay),
cropped to its icon, centred in a square with a small margin and scaled down without fringes,
the alpha made hard.

Usage:
  cut_icons.py raw.png out_prefix --grid ROWSxCOLS [--count N] [--margin 0.06]

Writes out_prefix-64.png and out_prefix-128.png: the icons side by side in reading order, and
out_prefix-preview.png (on a grey background, 2x). Needs Pillow.
"""
import argparse

from PIL import Image


def keyed(p):
    r, g, b = p[:3]
    return r > 170 and b > 170 and g < 90 and abs(r - b) < 70


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("raw")
    ap.add_argument("out")
    ap.add_argument("--grid", required=True)
    ap.add_argument("--count", type=int, default=0)
    ap.add_argument("--margin", type=float, default=0.06)
    a = ap.parse_args()
    rows, cols = (int(v) for v in a.grid.split("x"))
    n = a.count or rows * cols

    img = Image.open(a.raw).convert("RGBA")
    px = img.load()
    for y in range(img.height):
        for x in range(img.width):
            if keyed(px[x, y]):
                px[x, y] = (0, 0, 0, 0)
    icons = []
    for i in range(n):
        r, c = divmod(i, cols)
        cell = img.crop((c * img.width // cols, r * img.height // rows, (c + 1) * img.width // cols, (r + 1) * img.height // rows))
        box = cell.getchannel("A").point(lambda v: 255 if v > 0 else 0).getbbox() or (0, 0, cell.width, cell.height)
        icon = cell.crop(box)
        side = int(max(icon.size) * (1 + 2 * a.margin))
        sq = Image.new("RGBA", (side, side), (0, 0, 0, 0))
        sq.paste(icon, ((side - icon.width) // 2, (side - icon.height) // 2))
        icons.append(sq)
    for size in (64, 128):
        strip = Image.new("RGBA", (size * n, size), (0, 0, 0, 0))
        for k, sq in enumerate(icons):
            small = sq.resize((size, size), Image.BOX)
            sp = small.load()
            for y in range(size):
                for x in range(size):
                    r, g, b, al = sp[x, y]
                    if al < 110:
                        sp[x, y] = (0, 0, 0, 0)
                    else:
                        f = 255 / al
                        sp[x, y] = (min(255, int(r * f)), min(255, int(g * f)), min(255, int(b * f)), 255)
            strip.paste(small, (k * size, 0))
        strip.save(f"{a.out}-{size}.png")
    prev = Image.new("RGBA", (64 * n, 64), (80, 76, 84, 255))
    prev.alpha_composite(Image.open(f"{a.out}-64.png"))
    prev.resize((prev.width * 2, prev.height * 2), Image.NEAREST).save(f"{a.out}-preview.png")
    print(f"{a.out}: {n} icons")


if __name__ == "__main__":
    main()