studio/spritesheet.py

"""FrameStack -> sprite sheet PNG.

A sprite sheet is a single PNG containing N animation frames in a grid. This is
the idiomatic asset format for 2D game engines (Babylon.js `SpriteManager`,
Phaser, PixiJS) — no codec dependency, instant load, deterministic playback.

Convention: frames packed left-to-right, top-to-bottom. Default 4 columns gives
square-ish output for 16-frame stacks (4x4 = 16 cells).
"""
from __future__ import annotations
from pathlib import Path

import numpy as np
from PIL import Image as PILImage

from pixel_cursor import FrameStack


def frame_stack_to_sprite_sheet(
    stack: FrameStack,
    cols: int = 4,
    bg_rgba: tuple[int, int, int, int] = (0, 0, 0, 0),
) -> PILImage.Image:
    """Pack T frames into a (rows*cols)-cell grid PNG. RGBA output."""
    t, h, w, c = stack.frames.shape
    if c not in (3, 4):
        raise ValueError(f"expected 3 or 4 channels; got shape {stack.frames.shape}")
    rows = (t + cols - 1) // cols
    sheet_w = w * cols
    sheet_h = h * rows

    sheet = np.zeros((sheet_h, sheet_w, 4), dtype=np.uint8)
    sheet[:, :] = bg_rgba

    for idx in range(t):
        r, k = divmod(idx, cols)
        y0, x0 = r * h, k * w
        frame = stack.frames[idx]
        if c == 3:
            sheet[y0 : y0 + h, x0 : x0 + w, :3] = frame
            sheet[y0 : y0 + h, x0 : x0 + w, 3] = 255
        else:
            sheet[y0 : y0 + h, x0 : x0 + w, :] = frame

    return PILImage.fromarray(sheet, mode="RGBA")


def write_sprite_sheet(
    stack: FrameStack,
    out_path: Path | str,
    cols: int = 4,
    bg_rgba: tuple[int, int, int, int] = (0, 0, 0, 0),
) -> Path:
    out = Path(out_path)
    img = frame_stack_to_sprite_sheet(stack, cols=cols, bg_rgba=bg_rgba)
    img.save(out, optimize=True)
    return out


def sheet_metadata(stack: FrameStack, cols: int = 4) -> dict:
    """Return a manifest describing the sheet (cell size, count, fps).

    Game engines need this to play back. Babylon SpriteManager expects
    `cellWidth, cellHeight, capacity (total cells)`.
    """
    t, h, w, _ = stack.frames.shape
    rows = (t + cols - 1) // cols
    return {
        "frames": t,
        "cols": cols,
        "rows": rows,
        "cell_width": w,
        "cell_height": h,
        "sheet_width": w * cols,
        "sheet_height": h * rows,
        "fps": stack.fps,
        "loop_duration_ms": int(round(1000 * t / max(1, stack.fps))),
    }