#!/usr/bin/env python3
"""Cut a generated sheet of floor tiles (a grid on a magenta background) into single tiles for
the high-res atlas: every tile cropped to its square, magenta inside it (chipped corners, holes)
filled with a dark void, optionally brightened, scaled to the tile size and reduced to one
palette shared by all tiles of the sheet.

Usage:
  make_floor_tiles.py raw.png out_prefix --grid 2x3 --pick 1,2,3,5,6 [--size 128]
                      [--colors 48] [--bright 1.0]

--pick      which cells to keep, numbered row by row from 1
--bright    factor on the colours (above 1: lighter)

Writes out_prefix-<size>.png: the picked tiles side by side, and out_prefix-<size>-preview.png
(a 6x4 floor of them laid out at random, scaled up 2x). Needs Pillow.
"""
import argparse
import random

from PIL import Image, ImageEnhance

VOID = (22, 14, 18)


def magenta(p):
    r, g, b = p[:3]
    return r > g + 60 and b > g + 60


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("raw")
    ap.add_argument("out")
    ap.add_argument("--grid", default="2x3")
    ap.add_argument("--pick", default="1,2,3,4,5")
    ap.add_argument("--size", type=int, default=128)
    ap.add_argument("--colors", type=int, default=48)
    ap.add_argument("--bright", type=float, default=1.0)
    a = ap.parse_args()
    rows, cols = (int(v) for v in a.grid.split("x"))
    pick = [int(v) - 1 for v in a.pick.split(",")]

    img = Image.open(a.raw).convert("RGB")
    w, h = img.size
    px = img.load()
    tiles = []
    for i in pick:
        r, c = divmod(i, cols)
        x0, x1, y0, y1 = c * w // cols, (c + 1) * w // cols, r * h // rows, (r + 1) * h // rows
        # the tile's square: the columns and rows that are mostly not magenta
        xs = [x for x in range(x0, x1) if sum(not magenta(px[x, y]) for y in range(y0, y1, 4)) > (y1 - y0) / 8]
        ys = [y for y in range(y0, y1) if sum(not magenta(px[x, y]) for x in range(x0, x1, 4)) > (x1 - x0) / 8]
        tile = img.crop((min(xs), min(ys), max(xs) + 1, max(ys) + 1))
        tp = tile.load()
        for y in range(tile.height):
            for x in range(tile.width):
                if magenta(tp[x, y]):
                    tp[x, y] = VOID
        if a.bright != 1.0:
            tile = ImageEnhance.Brightness(tile).enhance(a.bright)
        tiles.append(tile.resize((a.size, a.size), Image.LANCZOS))

    strip = Image.new("RGB", (a.size * len(tiles), a.size))
    for k, t in enumerate(tiles):
        strip.paste(t, (k * a.size, 0))
    strip = strip.quantize(colors=a.colors, method=Image.Quantize.MEDIANCUT, dither=Image.Dither.NONE).convert("RGB")
    strip.save(f"{a.out}-{a.size}.png")

    rnd = random.Random(7)
    prev = Image.new("RGB", (a.size * 6, a.size * 4))
    for k in range(24):
        v = rnd.randrange(len(tiles))
        prev.paste(strip.crop((v * a.size, 0, (v + 1) * a.size, a.size)), ((k % 6) * a.size, (k // 6) * a.size))
    prev.resize((prev.width * 2, prev.height * 2), Image.NEAREST).save(f"{a.out}-{a.size}-preview.png")
    print(f"{a.out}-{a.size}.png: {len(tiles)} tiles")


if __name__ == "__main__":
    main()