nonogram_print.py

#!/usr/bin/env python3
"""Generate and print a nonogram (picross) puzzle on the S01 thermal printer.

Each row and column is labelled with the run-lengths of filled cells; solve them
together to reveal a hidden picture. Pictures are drawn programmatically (heart,
star, etc.) so every puzzle is clean and solvable. Like the maze and sudoku, the
seed prints on the card and --solution reveals the answer.

Examples:
    python nonogram_print.py                 # random picture
    python nonogram_print.py --shape star
    python nonogram_print.py --seed 1234
    python nonogram_print.py --seed 1234 --solution
    python nonogram_print.py --no-print
"""
from __future__ import annotations

import argparse
from functools import lru_cache
from math import cos, pi, sin
from pathlib import Path
import random

from PIL import Image, ImageDraw

from print_common import ROOT, Card, font

SHAPES = ["heart", "star", "diamond", "plus", "circle", "ring", "arrow", "smiley"]


def shape_image(name: str, s: int) -> Image.Image:
    im = Image.new("L", (s, s), 255)
    d = ImageDraw.Draw(im)
    cx = cy = s / 2
    if name == "circle":
        d.ellipse((s * .14, s * .14, s * .86, s * .86), fill=0)
    elif name == "ring":
        d.ellipse((s * .10, s * .10, s * .90, s * .90), fill=0)
        d.ellipse((s * .34, s * .34, s * .66, s * .66), fill=255)
    elif name == "diamond":
        d.polygon([(cx, s * .10), (s * .90, cy), (cx, s * .90), (s * .10, cy)], fill=0)
    elif name == "plus":
        d.rectangle((s * .40, s * .12, s * .60, s * .88), fill=0)
        d.rectangle((s * .12, s * .40, s * .88, s * .60), fill=0)
    elif name == "arrow":
        d.polygon([(cx, s * .10), (s * .86, cy), (s * .63, cy), (s * .63, s * .90),
                   (s * .37, s * .90), (s * .37, cy), (s * .14, cy)], fill=0)
    elif name == "heart":
        rr, top = s * 0.22, s * 0.34
        d.ellipse((cx - 2 * rr, top - rr, cx, top + rr), fill=0)
        d.ellipse((cx, top - rr, cx + 2 * rr, top + rr), fill=0)
        d.polygon([(cx - 2 * rr, top), (cx + 2 * rr, top), (cx, top + s * 0.44)], fill=0)
    elif name == "smiley":
        d.ellipse((s * .12, s * .12, s * .88, s * .88), fill=0)
        d.ellipse((s * .33, s * .35, s * .43, s * .47), fill=255)
        d.ellipse((s * .57, s * .35, s * .67, s * .47), fill=255)
        d.arc((s * .30, s * .40, s * .70, s * .75), 20, 160, fill=255, width=int(s * .05))
    else:  # star
        pts = []
        for i in range(10):
            ang = -pi / 2 + i * pi / 5
            rad = s * 0.45 if i % 2 == 0 else s * 0.19
            pts.append((cx + rad * cos(ang), cy + rad * sin(ang)))
        d.polygon(pts, fill=0)
    return im


def make_grid(name: str, n: int) -> list[list[int]]:
    s = n * 16
    im = shape_image(name, s)
    px = im.load()
    grid = [[0] * n for _ in range(n)]
    for r in range(n):
        for c in range(n):
            dark = sum(px[c * 16 + dx, r * 16 + dy] < 128 for dx in range(2, 14, 3) for dy in range(2, 14, 3))
            grid[r][c] = 1 if dark >= 5 else 0  # majority of sampled sub-cells filled
    return grid


def clues(line: list[int]) -> list[int]:
    runs, run = [], 0
    for v in line:
        if v:
            run += 1
        elif run:
            runs.append(run)
            run = 0
    if run:
        runs.append(run)
    return runs or [0]


@lru_cache(maxsize=None)
def _line_candidates(clue: tuple, length: int) -> list:
    """All filled/empty arrangements (1/0) of *length* matching *clue*."""
    if clue == (0,) or not clue:
        return [(0,) * length]
    first, rest = clue[0], clue[1:]
    rest_min = sum(rest) + len(rest)  # runs in rest each need a leading gap
    out = []
    for start in range(0, length - first - rest_min + 1):
        head = (0,) * start + (1,) * first
        if rest:
            for tail in _line_candidates(rest, length - start - first - 1):
                out.append(head + (0,) + tail)
        else:
            out.append(head + (0,) * (length - start - first))
    return out


def _deduce(known: tuple, clue: tuple) -> tuple:
    """Set cells that every clue-consistent candidate agrees on (1=fill, 2=empty)."""
    cands = [c for c in _line_candidates(clue, len(known))
             if all(k == 0 or (k == 1) == (v == 1) for k, v in zip(known, c))]
    new = list(known)
    for i in range(len(known)):
        vals = {c[i] for c in cands}
        if len(vals) == 1:
            new[i] = 1 if vals.pop() == 1 else 2
    return tuple(new)


def line_solvable(row_clues: list, col_clues: list, n: int) -> bool:
    """True if pure constraint propagation fully determines the grid (=> unique)."""
    grid = [[0] * n for _ in range(n)]
    changed = True
    while changed:
        changed = False
        for r in range(n):
            cur = tuple(grid[r])
            new = _deduce(cur, tuple(row_clues[r]))
            if new != cur:
                grid[r] = list(new)
                changed = True
        for c in range(n):
            cur = tuple(grid[r][c] for r in range(n))
            new = _deduce(cur, tuple(col_clues[c]))
            if new != cur:
                for r in range(n):
                    grid[r][c] = new[r]
                changed = True
    return all(cell != 0 for row in grid for cell in row)


def grid_clues(grid: list, n: int):
    rows = [clues(grid[r]) for r in range(n)]
    cols = [clues([grid[r][c] for r in range(n)]) for c in range(n)]
    return rows, cols


def render(grid, n, show_solution, target_w):
    rows = [clues(grid[r]) for r in range(n)]
    cols = [clues([grid[r][c] for r in range(n)]) for c in range(n)]
    cs = 15  # clue number spacing
    left = max(len(r) for r in rows) * cs + 6
    top = max(len(c) for c in cols) * cs + 6
    # size the cell so clues + grid fit the printable width, then pad so the
    # final right/bottom border lines aren't clipped at the image edge
    cell = min(22, (target_w - left - 4) // n)
    gpx = cell * n
    W, H = left + gpx + 3, top + gpx + 3
    img = Image.new("L", (W, H), 255)
    d = ImageDraw.Draw(img)
    clue_f = font(True, 13)

    if show_solution:
        for r in range(n):
            for c in range(n):
                if grid[r][c]:
                    x0, y0 = left + c * cell, top + r * cell
                    d.rectangle((x0 + 1, y0 + 1, x0 + cell - 1, y0 + cell - 1), fill=0)

    for i in range(n + 1):
        w = 3 if i % 5 == 0 else 1
        d.line((left + i * cell, top, left + i * cell, H), fill=0, width=w)
        d.line((left, top + i * cell, W, top + i * cell), fill=0, width=w)

    for r in range(n):
        y = top + r * cell + cell / 2
        for i, num in enumerate(reversed(rows[r])):
            d.text((left - 6 - i * cs, y), str(num), fill=0, font=clue_f, anchor="rm")
    for c in range(n):
        x = left + c * cell + cell / 2
        for i, num in enumerate(reversed(cols[c])):
            d.text((x, top - 6 - i * cs), str(num), fill=0, font=clue_f, anchor="mm")
    return img


def main() -> int:
    parser = argparse.ArgumentParser(description="Generate and print a nonogram on the S01 thermal printer.")
    parser.add_argument("--shape", choices=SHAPES, default=None, help="Hidden picture (default: random).")
    parser.add_argument("--size", type=int, default=None, help="Grid size n x n (default: random 10-14).")
    parser.add_argument("--seed", type=int, default=None, help="Seed for a reproducible puzzle.")
    parser.add_argument("--solution", action="store_true", help="Reveal the picture.")
    parser.add_argument("--out", type=Path, default=None)
    parser.add_argument("--darkness", type=int, choices=range(1, 6), default=3)
    parser.add_argument("--bottom-feed", type=int, default=24)
    parser.add_argument("--no-print", action="store_true")
    args = parser.parse_args()

    if args.size is not None and not 8 <= args.size <= 14:
        parser.error("--size must be 8-14")
    seed = args.seed if args.seed is not None else random.randrange(1, 100000)
    rng = random.Random(seed)
    n = args.size or rng.choice([10, 11, 12, 13, 14])

    # only keep a puzzle if it can be solved by pure logic (i.e. has a unique
    # solution); try shapes/sizes until one qualifies
    order = [args.shape] if args.shape else rng.sample(SHAPES, len(SHAPES))
    shape, grid, solvable = order[0], make_grid(order[0], n), False
    for s in order:
        for size in (n, n + 1, n - 1, n + 2):
            if not 8 <= size <= 15:
                continue
            g = make_grid(s, size)
            rows, cols = grid_clues(g, size)
            if line_solvable(rows, cols, size):
                shape, grid, n, solvable = s, g, size, True
                break
        if solvable:
            break
    note = "" if solvable else "  (needs guessing!)"
    print(f"Generating '{shape}' nonogram (seed {seed}, {n}x{n}){note} ...")
    args.size = n

    card = Card()
    card.title("NONOGRAM")
    label = f"ANSWER: {shape.upper()}" if args.solution else "fill cells to reveal the picture"
    card.line(label, size=12, bold=args.solution, center=True).gap(4)
    card.image(render(grid, args.size, args.solution, card.inner_w), border=False)
    card.footer("picross", f"seed {seed}")
    suffix = "_solved" if args.solution else ""
    out = args.out or ROOT / f"nonogram_{shape}_{seed}{suffix}.png"
    return card.finish(out, args.bottom_feed, args.darkness, do_print=not args.no_print)


if __name__ == "__main__":
    raise SystemExit(main())