#!/usr/bin/env python3
"""Crunch an image into a limited-palette pixel-art render.

Pipeline:  load -> resize-to-target (any dims/aspect, fill/fit/pad + focus)
           -> palette-constrained dither (OKLab-perceptual match, linear-light
              error diffusion) -> indexed PNG + upscaled preview.

Run with the sibling venv:  .venv/bin/python pixelate.py <image> [opts]

LOCKED Primordyn recipe (2026-06-01): the DEFAULTS are the look. Running
    pixelate.py render.png --palette primordyn_v2.json
is exactly  --size 640x360 --dither floyd --space oklab --mode fill  (serpentine,
strength 1.0). Override --size 320x180 for the chunkier cut.

The quality knobs that matter: perceptual matching in OKLab (right color, not
just RGB-nearest), error diffusion done in LINEAR light (gamma-correct, no
muddy darkening), and a choice of diffusion kernels (Floyd-Steinberg, Atkinson,
Jarvis, Stucki, Sierra) or ordered Bayer.

Palettes auto-detected: JSON ({"colors":[[r,g,b],...]} or bare list), GIMP .gpl,
JASC/Aseprite .pal, hex lists (.hex/.txt), Paint.NET, .png swatch, or adaptive:N.
"""
import argparse
import json
import os
import re
import sys

import numpy as np
from PIL import Image


# --------------------------------------------------------------------------- #
# Palette loading (format-agnostic)
# --------------------------------------------------------------------------- #
def _hex_to_rgb(h):
    h = h.lstrip("#")
    if h.lower().startswith("0x"):
        h = h[2:]
    if len(h) == 8:                 # 8 digits: Paint.NET AARRGGBB -> drop alpha
        h = h[2:]
    return (int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16))


def _parse_text_palette(text):
    lines = text.splitlines()
    first = next((ln.strip() for ln in lines if ln.strip()), "")

    if first.lower().startswith("gimp palette"):           # GIMP .gpl
        out = []
        for ln in lines[1:]:
            s = ln.strip()
            if not s or s.startswith("#") or s.lower().startswith(("name:", "columns:")):
                continue
            parts = s.split()
            if len(parts) >= 3 and all(p.isdigit() for p in parts[:3]):
                out.append(tuple(int(p) for p in parts[:3]))
        return out

    if first.upper().startswith("JASC-PAL"):               # JASC / Aseprite .pal
        return [tuple(int(p) for p in ln.split()[:3])
                for ln in lines[3:] if len(ln.split()) >= 3 and ln.split()[0].isdigit()]

    hexes = re.findall(r"(?:#|0x)?([0-9a-fA-F]{8}|[0-9a-fA-F]{6})\b", text)
    if hexes and ("#" in text or "0x" in text.lower() or any(re.search("[a-fA-F]", h) for h in hexes)):
        return [_hex_to_rgb(h) for h in hexes]

    out = []                                               # decimal triples
    for ln in lines:
        s = ln.strip()
        if not s or s.startswith((";", "//")):
            continue
        nums = [p for p in re.split(r"[\s,]+", s) if p.isdigit()]
        if len(nums) >= 3:
            out.append(tuple(int(p) for p in nums[:3]))
    return out


def load_palette(spec, ref_image=None):
    if spec.startswith("adaptive:"):
        n = int(spec.split(":", 1)[1])
        q = ref_image.convert("RGB").quantize(colors=n, method=Image.MEDIANCUT)
        return np.array(q.getpalette()[: n * 3], dtype=np.float32).reshape(-1, 3)

    raw = open(spec, "rb").read()
    if spec.lower().endswith(".json") or raw.lstrip()[:1] in (b"{", b"["):
        try:
            doc = json.loads(raw.decode("utf-8"))
            colors = doc.get("colors", doc) if isinstance(doc, dict) else doc
            triples = [tuple(c[:3]) for c in colors if isinstance(c, (list, tuple)) and len(c) >= 3]
            if triples:
                return np.array(triples, dtype=np.float32)
        except (ValueError, UnicodeDecodeError):
            pass

    if raw[:8] == b"\x89PNG\r\n\x1a\n":
        arr = np.array(Image.open(spec).convert("RGB")).reshape(-1, 3)
        return np.array(list(dict.fromkeys(map(tuple, arr.tolist()))), dtype=np.float32)

    colors = _parse_text_palette(raw.decode("utf-8", "replace"))
    if not colors:
        sys.exit(f"!! could not parse any colors from {spec}")
    return np.array(colors, dtype=np.float32)


# --------------------------------------------------------------------------- #
# Color spaces. Matching happens in a perceptual space; diffusion in linear.
# --------------------------------------------------------------------------- #
def srgb_to_linear(c):
    return np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)


def linear_to_srgb(c):
    return np.where(c <= 0.0031308, c * 12.92, 1.055 * np.clip(c, 0, None) ** (1 / 2.4) - 0.055)


def linear_to_oklab(c):
    r, g, b = c[..., 0], c[..., 1], c[..., 2]
    l = 0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b
    m = 0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b
    s = 0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b
    l_, m_, s_ = np.cbrt(l), np.cbrt(m), np.cbrt(s)
    return np.stack([
        0.2104542553 * l_ + 0.7936177850 * m_ - 0.0040720468 * s_,
        1.9779984951 * l_ - 2.4285922050 * m_ + 0.4505937099 * s_,
        0.0259040371 * l_ + 0.7827717662 * m_ - 0.8086757660 * s_,
    ], -1)


def to_match_fn(space):
    return {"oklab": linear_to_oklab, "linear": lambda c: c, "srgb": linear_to_srgb}[space]


# --------------------------------------------------------------------------- #
# Nearest-color LUT over the linear cube (built via the perceptual metric).
# Lets the sequential error-diffusion loop do O(1) lookups instead of an
# argmin-over-256 per pixel.
# --------------------------------------------------------------------------- #
def build_lut(palette_match, to_match, grid=64):
    axis = np.linspace(0.0, 1.0, grid, dtype=np.float32)
    gx, gy, gz = np.meshgrid(axis, axis, axis, indexing="ij")
    pts_lin = np.stack([gx, gy, gz], -1).reshape(-1, 3)
    pts = to_match(pts_lin)
    out = np.empty(pts.shape[0], dtype=np.int32)
    step = 4096
    for s in range(0, pts.shape[0], step):
        d = ((pts[s:s + step, None, :] - palette_match[None, :, :]) ** 2).sum(2)
        out[s:s + step] = d.argmin(1)
    return out.reshape(grid, grid, grid)


# kernel = (divisor, [(dx, dy, weight), ...]); error spread to *future* pixels.
KERNELS = {
    "floyd":    (16, [(1, 0, 7), (-1, 1, 3), (0, 1, 5), (1, 1, 1)]),
    "atkinson": (8,  [(1, 0, 1), (2, 0, 1), (-1, 1, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1)]),
    "jjn":      (48, [(1, 0, 7), (2, 0, 5), (-2, 1, 3), (-1, 1, 5), (0, 1, 7), (1, 1, 5),
                      (2, 1, 3), (-2, 2, 1), (-1, 2, 3), (0, 2, 5), (1, 2, 3), (2, 2, 1)]),
    "stucki":   (42, [(1, 0, 8), (2, 0, 4), (-2, 1, 2), (-1, 1, 4), (0, 1, 8), (1, 1, 4),
                      (2, 1, 2), (-2, 2, 1), (-1, 2, 2), (0, 2, 4), (1, 2, 2), (2, 2, 1)]),
    "sierra":   (32, [(1, 0, 5), (2, 0, 3), (-2, 1, 2), (-1, 1, 4), (0, 1, 5), (1, 1, 4),
                      (2, 1, 2), (-1, 2, 2), (0, 2, 3), (1, 2, 2)]),
}

BAYER = {
    4: np.array([[0, 8, 2, 10], [12, 4, 14, 6], [3, 11, 1, 9], [15, 7, 13, 5]], np.float32) / 16.0 - 0.5,
    8: (np.array([[0, 32, 8, 40, 2, 34, 10, 42], [48, 16, 56, 24, 50, 18, 58, 26],
                  [12, 44, 4, 36, 14, 46, 6, 38], [60, 28, 52, 20, 62, 30, 54, 22],
                  [3, 35, 11, 43, 1, 33, 9, 41], [51, 19, 59, 27, 49, 17, 57, 25],
                  [15, 47, 7, 39, 13, 45, 5, 37], [63, 31, 55, 23, 61, 29, 53, 21]], np.float32) / 64.0 - 0.5),
}


def _lut_lookup(lin_vals, lut, gscale):
    gi = np.clip((lin_vals * gscale + 0.5).astype(np.int32), 0, gscale)
    return lut[gi[..., 0], gi[..., 1], gi[..., 2]]


def error_diffuse(lin, lut, palette_lin, kernel, serpentine, strength):
    h, w, _ = lin.shape
    work = lin.copy()
    idxmap = np.zeros((h, w), np.int32)
    div, taps = kernel
    gscale = lut.shape[0] - 1
    for y in range(h):
        rev = serpentine and (y & 1)
        xr = range(w - 1, -1, -1) if rev else range(w)
        for x in xr:
            v = work[y, x]
            gi = np.clip((v * gscale + 0.5).astype(np.int32), 0, gscale)
            i = int(lut[gi[0], gi[1], gi[2]])
            idxmap[y, x] = i
            err = (v - palette_lin[i]) * strength
            for dx, dy, wt in taps:
                nx, ny = x + (-dx if rev else dx), y + dy
                if 0 <= nx < w and 0 <= ny < h:
                    work[ny, nx] += err * (wt / div)
    return idxmap


def ordered(lin, lut, bayer, strength):
    h, w, _ = lin.shape
    tile = np.tile(bayer, (h // bayer.shape[0] + 1, w // bayer.shape[1] + 1))[:h, :w]
    v = np.clip(lin + tile[..., None] * strength, 0.0, 1.0)
    return _lut_lookup(v, lut, lut.shape[0] - 1)


# --------------------------------------------------------------------------- #
# Resize to target dimensions (ported from convert_image.py's crop logic)
# --------------------------------------------------------------------------- #
def resize_to_target(img, tw, th, mode, fx, fy, resample):
    sw, sh = img.size
    sa, ta = sw / sh, tw / th
    if mode == "pad":
        scale = min(tw / sw, th / sh)
        nw, nh = max(1, round(sw * scale)), max(1, round(sh * scale))
        canvas = Image.new("RGB", (tw, th), (0, 0, 0))
        canvas.paste(img.resize((nw, nh), resample), ((tw - nw) // 2, (th - nh) // 2))
        return canvas
    if mode == "fit":                                  # scale to cover, crop minimal
        scale = max(tw / sw, th / sh)
        nw, nh = max(tw, round(sw * scale)), max(th, round(sh * scale))
        scaled = img.resize((nw, nh), resample)
        x, y = round((nw - tw) * fx), round((nh - th) * fy)
        return scaled.crop((x, y, x + tw, y + th))
    # fill (default): crop to target aspect, then scale
    if sa > ta:
        cw, ch = round(sh * ta), sh
    else:
        cw, ch = sw, round(sw / ta)
    x, y = round((sw - cw) * fx), round((sh - ch) * fy)
    return img.crop((x, y, x + cw, y + ch)).resize((tw, th), resample)


FOCI = {"center": (.5, .5), "top": (.5, 0), "bottom": (.5, 1), "left": (0, .5), "right": (1, .5),
        "top-left": (0, 0), "top-right": (1, 0), "bottom-left": (0, 1), "bottom-right": (1, 1)}
PRESETS = {"primordyn": (640, 360), "640p": (640, 360), "small": (320, 180), "320p": (320, 180),
           "vic20": (176, 184), "c64": (320, 200)}


def parse_size(s):
    if s.lower() in PRESETS:
        return PRESETS[s.lower()]
    m = re.fullmatch(r"(\d+)\s*[xX]\s*(\d+)", s.strip())
    if not m:
        sys.exit(f"!! --size must be WxH or one of {list(PRESETS)}; got '{s}'")
    return int(m.group(1)), int(m.group(2))


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("image")
    p.add_argument("--palette", default="primordyn_v2.json", help="palette file or adaptive:N")
    p.add_argument("--size", default="640x360", help=f"WxH or preset {list(PRESETS)}")
    p.add_argument("--mode", choices=["fill", "fit", "pad"], default="fill", help="aspect handling")
    p.add_argument("--focus", default="center", help=f"crop focus: {list(FOCI)} or 'fx,fy' (0..1)")
    p.add_argument("--dither", choices=list(KERNELS) + ["bayer4", "bayer8", "none"], default="floyd")
    p.add_argument("--space", choices=["oklab", "linear", "srgb"], default="oklab", help="color-match space")
    p.add_argument("--strength", type=float, default=1.0, help="dither amount (error-diffusion fraction / bayer scale)")
    p.add_argument("--no-serpentine", action="store_true")
    p.add_argument("--resample", choices=["lanczos", "box", "bilinear", "hamming", "nearest"], default="lanczos")
    p.add_argument("--lut", type=int, default=64, help="nearest-color LUT resolution per axis")
    p.add_argument("--preview-scale", type=int, default=2, help="nearest-neighbor upscale for the preview (0=off)")
    p.add_argument("--out", default=None)
    args = p.parse_args()

    tw, th = parse_size(args.size)
    fx, fy = FOCI.get(args.focus.lower(), None) or tuple(float(v) for v in args.focus.split(","))
    rfilt = {"lanczos": Image.LANCZOS, "box": Image.BOX, "bilinear": Image.BILINEAR,
             "hamming": Image.HAMMING, "nearest": Image.NEAREST}[args.resample]

    src = Image.open(args.image).convert("RGB")
    small = resize_to_target(src, tw, th, args.mode, fx, fy, rfilt)

    palette_srgb = load_palette(args.palette, ref_image=small)
    pal01 = palette_srgb / 255.0
    palette_lin = srgb_to_linear(pal01).astype(np.float32)
    to_match = to_match_fn(args.space)
    palette_match = to_match(palette_lin).astype(np.float32)

    lin = srgb_to_linear(np.asarray(small, np.float32) / 255.0)

    if args.dither in KERNELS:
        lut = build_lut(palette_match, to_match, args.lut)
        idxmap = error_diffuse(lin, lut, palette_lin, KERNELS[args.dither],
                               not args.no_serpentine, args.strength)
    elif args.dither.startswith("bayer"):
        lut = build_lut(palette_match, to_match, args.lut)
        idxmap = ordered(lin, lut, BAYER[int(args.dither[5:])], args.strength)
    else:  # none -> hard nearest, vectorized
        flat = to_match(lin).reshape(-1, 3)
        out = np.empty(flat.shape[0], np.int32)
        for s in range(0, flat.shape[0], 65536):
            d = ((flat[s:s + 65536, None, :] - palette_match[None, :, :]) ** 2).sum(2)
            out[s:s + 65536] = d.argmin(1)
        idxmap = out.reshape(th, tw)

    # --- outputs: true indexed PNG (target res) + upscaled RGB preview ---
    base = os.path.splitext(args.out or args.image)[0]
    lo_path = args.out or f"{base}_{tw}x{th}_{args.dither}.png"
    idx_img = Image.frombytes("P", (tw, th), idxmap.astype(np.uint8).tobytes())
    flat_pal = palette_srgb.astype(np.uint8).reshape(-1).tolist()
    idx_img.putpalette(flat_pal + [0] * (768 - len(flat_pal)))
    idx_img.save(lo_path)

    used = len(np.unique(idxmap))
    print(f"palette: {len(palette_srgb)} colors ({args.palette}) | used: {used}")
    print(f"target: {tw}x{th} via {args.mode}/{args.resample} focus={args.focus}")
    print(f"dither: {args.dither} | match-space: {args.space} | strength: {args.strength}"
          f"{'' if args.no_serpentine or args.dither in ('none',) or args.dither.startswith('bayer') else ' | serpentine'}")
    print(f"wrote {lo_path} (indexed PNG)")
    if args.preview_scale > 1:
        up_path = os.path.splitext(lo_path)[0] + f"_x{args.preview_scale}.png"
        idx_img.convert("RGB").resize((tw * args.preview_scale, th * args.preview_scale),
                                      Image.NEAREST).save(up_path)
        print(f"wrote {up_path} (preview)")


if __name__ == "__main__":
    main()