"""
AI "fake pixel art" -> real Minecraft item texture.

    python art/unfake_texture.py art/emerald_sword_source.png art/emerald_sword
    python art/unfake_texture.py src.png out --outline keep --method dominant --colors32 16

Writes <out>_snapped.png (1 px per detected fake pixel, bg removed, cropped, outline recolored),
<out>_32.png and <out>_16.png (quantized), plus <out>_preview.png.

Steps: detect the fake-pixel grid (edge autocorrelation); snap each cell to its median colour; flood-fill the
background from the border; crop; recolor the AI's black outline to a darkened copy of the material next to it
(vanilla items use material-dark edges, never black); pad so the pommel sits bottom-left and the tip top-right;
area-average downscale (alpha-weighted, coverage threshold); palette quantize.
Needs: pip install pillow numpy
"""
import argparse
from collections import Counter, deque

import numpy as np
from PIL import Image

BG_TOL = 28          # background colour tolerance (per channel)
CELL_INSET = 4       # ignore this many px at each cell edge when sampling (AI cells bleed)


def detect_grid(rgb):
    g = rgb.astype(np.float32).mean(axis=2)
    dx = np.abs(np.diff(g, axis=1)).sum(axis=0)
    dy = np.abs(np.diff(g, axis=0)).sum(axis=1)

    def period(sig, lo=4, hi=96):
        s = sig - sig.mean()
        ac = np.correlate(s, s, mode="full")[len(s) - 1:]
        ac /= ac[0]
        return max(range(lo, hi), key=lambda k: ac[k])

    p = round((period(dx) + period(dy)) / 2)
    ox = max(range(p), key=lambda o: dx[o::p].sum()) + 1
    oy = max(range(p), key=lambda o: dy[o::p].sum()) + 1
    return p, ox, oy


def snap(rgb, p, ox, oy):
    h, w = rgb.shape[:2]
    nx, ny = (w - ox) // p, (h - oy) // p
    out = np.zeros((ny, nx, 3), np.uint8)
    i0, i1 = CELL_INSET, p - CELL_INSET
    for j in range(ny):
        for i in range(nx):
            cell = rgb[oy + j * p + i0: oy + j * p + i1, ox + i * p + i0: ox + i * p + i1].reshape(-1, 3)
            out[j, i] = np.median(cell, axis=0)
    return out


def strip_background(img):
    ny, nx = img.shape[:2]
    corners = np.array([img[0, 0], img[0, -1], img[-1, 0], img[-1, -1]], np.int32)
    bg = np.median(corners, axis=0)
    bg_like = np.abs(img.astype(np.int32) - bg).max(axis=2) <= BG_TOL
    mask = np.zeros((ny, nx), bool)
    q = deque()
    for j in range(ny):
        for i in range(nx):
            if (j in (0, ny - 1) or i in (0, nx - 1)) and bg_like[j, i]:
                mask[j, i] = True
                q.append((j, i))
    while q:
        j, i = q.popleft()
        for jj, ii in ((j + 1, i), (j - 1, i), (j, i + 1), (j, i - 1)):
            if 0 <= jj < ny and 0 <= ii < nx and not mask[jj, ii] and bg_like[jj, ii]:
                mask[jj, ii] = True
                q.append((jj, ii))
    rgba = np.dstack([img, np.where(mask, 0, 255).astype(np.uint8)])
    return crop(rgba)


def crop(rgba):
    ys, xs = np.where(rgba[:, :, 3] > 0)
    return rgba[ys.min(): ys.max() + 1, xs.min(): xs.max() + 1]


def touches_transparent(a, y, x):
    h, w = a.shape[:2]
    for dy in (-1, 0, 1):
        for dx in (-1, 0, 1):
            if dy == 0 and dx == 0:
                continue
            yy, xx = y + dy, x + dx
            if yy < 0 or yy >= h or xx < 0 or xx >= w or a[yy, xx, 3] == 0:
                return True
    return False


def recolor_outline(a, dark=40, factor=0.55, passes=2, tint=None, tint_mix=0.5):
    """Near-black pixels on the silhouette become a darkened copy of the nearest real colour next to them,
    optionally blended toward `tint` (an (r,g,b)) so the whole edge reads as one material."""
    a = a.copy()
    h, w = a.shape[:2]
    tint = None if tint is None else np.array(tint, np.float32)
    for _ in range(passes):
        ol = np.zeros((h, w), bool)
        for y in range(h):
            for x in range(w):
                ol[y, x] = bool(a[y, x, 3]) and a[y, x, :3].max() < dark and touches_transparent(a, y, x)
        if not ol.any():
            break
        new = a.copy()
        for y, x in zip(*np.where(ol)):
            cols = []
            for r in (1, 2):
                for dy in range(-r, r + 1):
                    for dx in range(-r, r + 1):
                        yy, xx = y + dy, x + dx
                        if 0 <= yy < h and 0 <= xx < w and a[yy, xx, 3] and not ol[yy, xx] and a[yy, xx, :3].max() >= dark:
                            cols.append(a[yy, xx, :3].astype(np.float32))
                if cols:
                    break
            if cols:
                c = np.mean(cols, axis=0) * factor
                if tint is not None:
                    c = c * (1 - tint_mix) + tint * factor * tint_mix
                new[y, x, :3] = np.clip(c, 0, 255).astype(np.uint8)
            else:
                new[y, x] = 0
        a = new
    return a


# Palette locks (like PixelBanana's palette option). Colours lifted from vanilla item textures.
PALETTES = {
    "vanilla-emerald": (
        "#002d00,#005300,#007b18,#009529,#00aa2c,#17dd62,#41f384,#82f6ad,#affdcd,#dbffeb"   # emerald.png
        ",#181818,#444444,#6b6b6b,#969696,#bebebe,#d8d8d8,#ffffff"                          # iron_sword.png grays
        ",#281e0b,#493615,#684e1e,#896727"                                                  # stick.png browns
    ),
}


def parse_palette(spec):
    if spec in PALETTES:
        spec = PALETTES[spec]
    elif not spec.startswith("#"):
        spec = ",".join(l.strip() for l in open(spec, encoding="utf-8") if l.strip())
    return np.array([[int(h[i:i + 2], 16) for i in (1, 3, 5)] for h in spec.split(",")], np.float32)


def apply_palette(img, pal):
    """Snap every opaque pixel to the nearest palette colour (Lab-ish weighting on RGB)."""
    out = img.copy()
    wgt = np.array([0.3, 0.59, 0.11], np.float32) * 3
    for y in range(img.shape[0]):
        for x in range(img.shape[1]):
            if img[y, x, 3]:
                d = (((pal - img[y, x, :3].astype(np.float32)) ** 2) * wgt).sum(1)
                out[y, x, :3] = pal[d.argmin()].astype(np.uint8)
    return out


def strip_outline(a, dark=35):
    a = a.copy()
    rm = [(y, x) for y in range(a.shape[0]) for x in range(a.shape[1])
          if a[y, x, 3] and a[y, x, :3].max() < dark and touches_transparent(a, y, x)]
    for y, x in rm:
        a[y, x] = 0
    return crop(a)


def pad_square(img, n):
    h, w = img.shape[:2]
    out = np.zeros((n, n, 4), np.uint8)
    out[n - h:, :w] = img          # pommel bottom-left, tip top-right (Minecraft's sword convention)
    return out


def down_mean(img, k, cover=0.5):
    """Alpha-weighted box filter; a block is opaque when at least `cover` of it was."""
    n = img.shape[0] // k
    out = np.zeros((n, n, 4), np.uint8)
    f = img.astype(np.float32)
    for y in range(n):
        for x in range(n):
            blk = f[y * k:(y + 1) * k, x * k:(x + 1) * k].reshape(-1, 4)
            a = blk[:, 3] / 255.0
            if a.mean() < cover:
                continue
            rgb = (blk[:, :3] * a[:, None]).sum(0) / a.sum()
            out[y, x] = (*np.clip(rgb, 0, 255).astype(np.uint8), 255)
    return out


def down_dominant(img, k):
    n = img.shape[0] // k
    out = np.zeros((n, n, 4), np.uint8)
    for y in range(n):
        for x in range(n):
            blk = img[y * k:(y + 1) * k, x * k:(x + 1) * k].reshape(-1, 4)
            op = blk[blk[:, 3] > 0]
            if len(op) * 2 < len(blk):
                continue
            best, cnt = Counter(map(tuple, op[:, :3])).most_common(1)[0]
            if cnt == 1:
                best = tuple(int(v) for v in np.median(op[:, :3], axis=0))
            out[y, x] = (*best, 255)
    return out


def quantize(img, ncol):
    q = Image.fromarray(img, "RGBA").quantize(colors=ncol, method=Image.Quantize.FASTOCTREE, dither=Image.Dither.NONE)
    a = np.asarray(q.convert("RGBA")).copy()
    a[img[:, :, 3] == 0] = 0
    return a


def preview(imgs, path, k):
    tiles = [np.asarray(Image.fromarray(i, "RGBA").resize((i.shape[1] * k, i.shape[0] * k), Image.NEAREST)) for i in imgs]
    h = max(t.shape[0] for t in tiles)
    tiles = [np.pad(t, ((0, h - t.shape[0]), (0, 0), (0, 0))) for t in tiles]
    sheet = np.concatenate(tiles, axis=1)
    bg = np.full(sheet.shape[:2] + (3,), 50, np.uint8)
    a = sheet[:, :, 3:4] / 255.0
    Image.fromarray((sheet[:, :, :3] * a + bg * (1 - a)).astype(np.uint8)).save(path)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("source")
    ap.add_argument("out", help="output prefix, e.g. art/emerald_sword")
    ap.add_argument("--outline", choices=["recolor", "keep", "strip"], default="recolor")
    ap.add_argument("--outline-factor", type=float, default=0.55, help="brightness of recolored outline vs neighbour")
    ap.add_argument("--outline-tint", default=None, help="hex colour the outline leans toward, e.g. #0a4a22")
    ap.add_argument("--tint-mix", type=float, default=0.5, help="0 = neighbour colour only, 1 = tint only")
    ap.add_argument("--method", choices=["mean", "dominant"], default="mean")
    ap.add_argument("--colors32", type=int, default=20)
    ap.add_argument("--colors16", type=int, default=12)
    ap.add_argument("--palette", default=None, help="palette lock: a name (%s), '#hex,#hex,...', or a file" % ", ".join(PALETTES))
    args = ap.parse_args()

    rgb = np.asarray(Image.open(args.source).convert("RGB"))
    p, ox, oy = detect_grid(rgb)
    sprite = strip_background(snap(rgb, p, ox, oy))
    print(f"grid: {p}px cells, offset ({ox},{oy}); sprite {sprite.shape[1]}x{sprite.shape[0]}")
    tint = None if args.outline_tint is None else tuple(int(args.outline_tint[i:i + 2], 16) for i in (1, 3, 5))
    if args.outline == "recolor":
        sprite = recolor_outline(sprite, factor=args.outline_factor, tint=tint, tint_mix=args.tint_mix)
    elif args.outline == "strip":
        sprite = strip_outline(sprite)
    pal = None if args.palette is None else parse_palette(args.palette)
    if pal is not None:
        sprite = apply_palette(sprite, pal)
    Image.fromarray(sprite, "RGBA").save(f"{args.out}_snapped.png")

    base = 64 if max(sprite.shape[:2]) <= 64 else 128
    padded = pad_square(sprite, base)
    Image.fromarray(padded, "RGBA").save(f"{args.out}_{base}.png")      # full-res: no downscale at all
    down = down_mean if args.method == "mean" else down_dominant
    s32 = down(padded, base // 32)
    s16 = down(padded, base // 16)
    s32 = apply_palette(s32, pal) if pal is not None else quantize(s32, args.colors32)
    s16 = apply_palette(s16, pal) if pal is not None else quantize(s16, args.colors16)
    Image.fromarray(s32, "RGBA").save(f"{args.out}_32.png")
    Image.fromarray(s16, "RGBA").save(f"{args.out}_16.png")
    preview([sprite, s32, s16], f"{args.out}_preview.png", 6)
    for name, img in (("32", s32), ("16", s16)):
        op = img[img[:, :, 3] > 0]
        print(f"{name}x{name}: {len(set(map(tuple, op[:, :3])))} colours, {len(op)} opaque px")


if __name__ == "__main__":
    main()
