#!/usr/bin/env python3
# Part of JoeyLib, Copyright (c) 2026 Scott Duensing, under the MIT License.
# See LICENSE.md for license information.
# SPDX-License-Identifier: MIT

"""traceLogo.py -- trace a flat-colour PNG logo into a layered SVG.

The JoeyLib logo ships as a 320x200 PNG, which is fine on screen and poor
everywhere else: the manual's cover wants it large, and a bitmap that small
either pixelates or blurs when scaled up. This produces a resolution-free SVG
of the same artwork so it can be set at any size.

How it works. The art is flat cartoon colour with antialiased edges, so:

  1. k-means the opaque pixels down to `--colors` clusters, which recovers the
     handful of real colours and throws away the antialiasing gradient,
  2. build a one-bit mask per colour,
  3. run each mask through potrace (the pure-Python `potracer`), which fits
     Bezier curves to the region outlines,
  4. emit one SVG <path> per colour.

Layer order matters: the paths are written largest-area first so big shapes sit
behind small ones, EXCEPT the darkest colour, which is written last. In cartoon
art the darkest colour is the ink outline, and it has to be on top or the fills
eat it.

Transparent pixels are left out of every mask, so the SVG keeps the PNG's
transparent background.

Needs `potracer`, `pillow` and `numpy`. They are not build dependencies -- the
checked-in SVG is the deliverable and this script only needs running if the
source art changes:

    python3 -m venv /tmp/tracevenv
    /tmp/tracevenv/bin/pip install potracer pillow numpy
    /tmp/tracevenv/bin/python tools/traceLogo.py assets/logo.png assets/logo.svg
"""

import argparse
import sys

import numpy as np
import potrace
from PIL import Image


def kmeans(samples, k, iterations=24, seed=1):
    """Tiny k-means over RGB samples. Returns (centres, labels)."""
    rng = np.random.default_rng(seed)
    # Seed the centres on well-separated samples (k-means++ style, simplified).
    centres = [samples[rng.integers(len(samples))].astype(np.float64)]
    for _ in range(k - 1):
        d = np.min([np.sum((samples - c) ** 2, axis=1) for c in centres], axis=0)
        total = d.sum()
        if total <= 0:
            centres.append(samples[rng.integers(len(samples))].astype(np.float64))
            continue
        centres.append(samples[np.searchsorted(np.cumsum(d / total),
                                               rng.random())].astype(np.float64))
    centres = np.array(centres)

    labels = np.zeros(len(samples), dtype=np.int32)
    for _ in range(iterations):
        dist = ((samples[:, None, :] - centres[None, :, :]) ** 2).sum(axis=2)
        newLabels = np.argmin(dist, axis=1)
        if np.array_equal(newLabels, labels):
            break
        labels = newLabels
        for i in range(k):
            hit = samples[labels == i]
            if len(hit):
                centres[i] = hit.mean(axis=0)
    return centres, labels


def pathData(path):
    """potrace path -> SVG 'd'. potracer already emits image-space
    coordinates (y increasing downward), so no axis flip is applied -- adding
    one renders the art upside down."""
    out = []
    for curve in path:
        start = curve.start_point
        out.append(f"M{start.x:.2f},{start.y:.2f}")
        for seg in curve:
            if seg.is_corner:
                out.append(f"L{seg.c.x:.2f},{seg.c.y:.2f}")
                out.append(f"L{seg.end_point.x:.2f},{seg.end_point.y:.2f}")
            else:
                out.append(
                    f"C{seg.c1.x:.2f},{seg.c1.y:.2f}"
                    f" {seg.c2.x:.2f},{seg.c2.y:.2f}"
                    f" {seg.end_point.x:.2f},{seg.end_point.y:.2f}")
        out.append("Z")
    return "".join(out)


def main(argv):
    ap = argparse.ArgumentParser(description="Trace a flat-colour PNG into SVG.")
    ap.add_argument("input")
    ap.add_argument("output")
    ap.add_argument("--colors", type=int, default=5,
                    help="number of flat colours to recover (default 5)")
    ap.add_argument("--alpha", type=int, default=128,
                    help="alpha above which a pixel counts as opaque (default 128)")
    ap.add_argument("--turdsize", type=int, default=2,
                    help="drop traced specks smaller than this many pixels")
    args = ap.parse_args(argv[1:])

    im = Image.open(args.input).convert("RGBA")
    width, height = im.size
    data = np.array(im)
    rgb = data[:, :, :3].astype(np.float64)
    opaque = data[:, :, 3] > args.alpha
    if not opaque.any():
        print(f"traceLogo: {args.input} is fully transparent", file=sys.stderr)
        return 1

    centres, labels = kmeans(rgb[opaque], args.colors)

    # Rebuild a full-size label image; -1 means transparent.
    labelImage = np.full((height, width), -1, dtype=np.int32)
    labelImage[opaque] = labels

    layers = []
    for i in range(args.colors):
        mask = labelImage == i
        area = int(mask.sum())
        if area == 0:
            continue
        # Two polarity traps, both measured against this logo:
        #   * pass the BOOLEAN mask, not 0/1 bytes -- potracer compares against
        #     blacklevel=0.5, so a 0/1 uint8 array reads as entirely
        #     foreground and traces the whole canvas as one rectangle;
        #   * pass the INVERSE -- potracer treats False as the ink it traces,
        #     so handing it the mask directly traces the background and every
        #     layer comes out as a full-canvas block.
        bitmap = potrace.Bitmap(~mask)
        traced = bitmap.trace(turdsize=args.turdsize)
        r, g, b = (int(round(v)) for v in centres[i])
        layers.append({
            "area": area,
            "color": f"#{r:02X}{g:02X}{b:02X}",
            "luma": 0.299 * r + 0.587 * g + 0.114 * b,
            "d": pathData(traced),
        })

    if not layers:
        print("traceLogo: nothing to trace", file=sys.stderr)
        return 1

    # Largest first so big shapes sit behind, then force the ink (darkest)
    # layer to the very top.
    layers.sort(key=lambda l: -l["area"])
    ink = min(layers, key=lambda l: l["luma"])
    layers.remove(ink)
    layers.append(ink)

    with open(args.output, "w") as f:
        f.write('<?xml version="1.0" encoding="UTF-8"?>\n')
        f.write(f'<svg xmlns="http://www.w3.org/2000/svg" '
                f'viewBox="0 0 {width} {height}" width="{width}" '
                f'height="{height}">\n')
        f.write(f'  <title>{args.input}</title>\n')
        for layer in layers:
            f.write(f'  <path fill="{layer["color"]}" '
                    f'fill-rule="evenodd" d="{layer["d"]}"/>\n')
        f.write('</svg>\n')

    print(f"traceLogo: {args.output} -- {len(layers)} colour layers, "
          f"{width}x{height} viewBox")
    for layer in sorted(layers, key=lambda l: -l["area"]):
        print(f"    {layer['color']}  {layer['area']:6d} px")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))