#!/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] [--inset 0.0] [--flatten]

--pick      which cells to keep, numbered row by row from 1
--colors    palette size (0: keep the colours; small coloured parts such as inlays would
            fade in a palette taken mostly by large grey areas)
--bright    factor on the colours (above 1: lighter)
--inset     cut this share off every side (a frame or edge joint the model drew around the tile)
--flatten   even out the light: no vignette towards the edges, all tiles equally bright, so
            the grid of tiles does not show on the floor

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, ImageFilter, ImageStat

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)
    ap.add_argument("--inset", type=float, default=0.0)
    ap.add_argument("--flatten", action="store_true")
    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.inset > 0:
            d = int(tile.width * a.inset)
            tile = tile.crop((d, d, tile.width - d, tile.height - d))
        tiles.append(tile)
    if a.flatten:
        # divide by the large-scale brightness: every tile gets the same, even brightness (the
        # sheet's mean), so neither a vignette nor a lighter or darker tile shows the grid
        blurs = [t.convert("L").filter(ImageFilter.GaussianBlur(t.width / 5)) for t in tiles]
        mean = sum(ImageStat.Stat(b).mean[0] for b in blurs) / len(blurs)
        for t, blur in zip(tiles, blurs):
            bp, tp = blur.load(), t.load()
            for y in range(t.height):
                for x in range(t.width):
                    f = mean / max(bp[x, y], 1)
                    tp[x, y] = tuple(min(255, int(c * f)) for c in tp[x, y])
    if a.bright != 1.0:
        tiles = [ImageEnhance.Brightness(t).enhance(a.bright) for t in tiles]
    tiles = [t.resize((a.size, a.size), Image.LANCZOS) for t in tiles]

    strip = Image.new("RGB", (a.size * len(tiles), a.size))
    for k, t in enumerate(tiles):
        strip.paste(t, (k * a.size, 0))
    if a.colors > 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()