"""Turn a MiniMax-H3 clip into game-ready sprites.

Everything here is CPU work on numpy arrays — the Space itself owns no GPU. The pipeline is the second half of
Victor Mustar's workflow: key out the flat backdrop, cut the clip down to its pose extremes, pack the extremes into
one sheet.

Keying happens in the **chroma plane of YCbCr**, not in RGB. Two properties matter and both fall out of dropping
luma:

* The contact shadow a text-to-image model insists on drawing under the character is the backdrop hue at a lower
  luminance. In the (Cb, Cr) plane it sits on top of the backdrop and keys out with it, for free.
* Black outlines and white gloves — the two things a naive RGB distance destroys — are *neutral*: they sit at the
  origin of the chroma plane, maximally far from a saturated backdrop. A rubber-hose character is mostly those two
  colors, so this is not a detail.

The backdrop color is measured off the frame border rather than assumed to be ``#FF00FF``. No text-to-image model
actually returns pure magenta when asked — the samples this was built against came back anywhere from ``#F0417F``
rose to ``#E845D8`` orchid — and keying against the requested color instead of the delivered one is the single
easiest way to get a matte full of holes.
"""

from __future__ import annotations

import base64
import json
import math
import os
from dataclasses import dataclass, field

import numpy as np
from PIL import Image, ImageFilter

# ---------------------------------------------------------------------------
# color
# ---------------------------------------------------------------------------

# ITU-R BT.601, the matrix JPEG/MPEG chroma subsampling is defined against — which is what the clip was encoded with.
_R2Y = np.array(
    [
        [0.299, 0.587, 0.114],
        [-0.168736, -0.331264, 0.5],
        [0.5, -0.418688, -0.081312],
    ],
    dtype=np.float32,
)
_Y2R = np.linalg.inv(_R2Y).astype(np.float32)


def to_ycc(rgb: np.ndarray) -> np.ndarray:
    """``(..., 3)`` float RGB in 0..1 -> YCbCr with luma in 0..1 and chroma in -0.5..0.5."""
    return rgb @ _R2Y.T


def to_rgb(ycc: np.ndarray) -> np.ndarray:
    return np.clip(ycc @ _Y2R.T, 0.0, 1.0)


def detect_key_color(frame: np.ndarray, border: float = 0.06) -> np.ndarray:
    """Measure the backdrop off the frame border, as float RGB in 0..1.

    The median over a ring — not the mean, and not a corner sample. A character whose foot or tail crosses into the
    ring shifts a mean; it has to cover more than half the ring's pixels to shift the median, which no centered
    sprite does.
    """
    h, w = frame.shape[:2]
    bh, bw = max(1, int(h * border)), max(1, int(w * border))
    ring = np.concatenate(
        [
            frame[:bh].reshape(-1, 3),
            frame[-bh:].reshape(-1, 3),
            frame[:, :bw].reshape(-1, 3),
            frame[:, -bw:].reshape(-1, 3),
        ]
    )
    return np.median(ring, axis=0)


# ---------------------------------------------------------------------------
# keying
# ---------------------------------------------------------------------------


def _smoothstep(edge0: float, edge1: float, x: np.ndarray) -> np.ndarray:
    if edge1 <= edge0:
        return (x > edge0).astype(np.float32)
    t = np.clip((x - edge0) / (edge1 - edge0), 0.0, 1.0)
    return t * t * (3.0 - 2.0 * t)


def matte(
    frame: np.ndarray,
    key_rgb: np.ndarray,
    tolerance: float = 0.10,
    softness: float = 0.09,
) -> np.ndarray:
    """The alpha channel alone, from the chroma distance to ``key_rgb``.

    ``tolerance`` is the chroma radius that is fully transparent and ``softness`` the width of the ramp out to fully
    opaque, both as distances in the (Cb, Cr) plane whose useful range is about 0..0.7.
    """
    ycc = to_ycc(frame)
    key_chroma = to_ycc(key_rgb[None, None, :])[0, 0][1:]
    if float(np.linalg.norm(key_chroma)) < 1e-4:
        # A neutral backdrop has no chroma to key against; nothing here can separate it from a white glove.
        return np.ones(frame.shape[:2], np.float32)
    delta = ycc[..., 1:] - key_chroma
    return _smoothstep(tolerance, tolerance + max(softness, 1e-4), np.linalg.norm(delta, axis=-1))


def refine_matte(alpha: np.ndarray, despeckle: int = 2, shrink: float = 0.2, feather: float = 0.6) -> np.ndarray:
    """Despeckle, choke and soften the matte — all of it before any color is recovered.

    A median filter removes single-pixel noise without eating a whisker, which an erode/dilate pair of the same
    radius would not.

    ``shrink`` is the compositor's choke, and it is what actually kills the halo. The alpha ramp is derived from
    chroma, and the clip arrives 4:2:0 subsampled, so chroma bleeds about a pixel past the real edge and a rim of
    genuinely-backdrop pixels comes out faintly opaque. Rescaling the ramp so everything under ``shrink`` goes to
    zero drops that rim.

    Ordering matters more than any of the individual filters here: feathering *expands* alpha outward, and if it runs
    after the color step it hands partial opacity to pixels that are still pure backdrop — a pink rim, reintroduced
    by the very filter meant to tidy the edge. The matte is finished here and `unmix` is told about it.
    """
    if despeckle > 0:
        image = Image.fromarray((np.clip(alpha, 0, 1) * 255).astype(np.uint8), "L")
        alpha = np.asarray(image.filter(ImageFilter.MedianFilter(size=despeckle * 2 + 1)), np.float32) / 255.0
    if shrink > 0:
        alpha = np.clip((alpha - shrink) / max(1.0 - shrink, 1e-3), 0.0, 1.0)
    if feather > 0:
        image = Image.fromarray((np.clip(alpha, 0, 1) * 255).astype(np.uint8), "L")
        alpha = np.asarray(image.filter(ImageFilter.GaussianBlur(radius=feather)), np.float32) / 255.0
    return alpha.astype(np.float32)


def _box_blur(x: np.ndarray, radius: int) -> np.ndarray:
    """Separable box blur over the first two axes, as a difference of prefix sums. No SciPy needed."""
    if radius <= 0:
        return x
    k = 2 * radius + 1
    out = x.astype(np.float32)
    for axis in (0, 1):
        pad = [(0, 0)] * out.ndim
        pad[axis] = (radius, radius)
        cumulative = np.cumsum(np.pad(out, pad, mode="edge"), axis=axis)
        zeros = list(cumulative.shape)
        zeros[axis] = 1
        cumulative = np.concatenate([np.zeros(zeros, np.float32), cumulative], axis=axis)
        size = cumulative.shape[axis]
        upper = np.take(cumulative, range(k, size), axis=axis)
        lower = np.take(cumulative, range(0, size - k), axis=axis)
        out = (upper - lower) / k
    return out


def extend_edges(rgb: np.ndarray, alpha: np.ndarray, radius: int = 3) -> np.ndarray:
    """Push opaque color outward into the matte's transition band.

    This is what finally removes the rim, and it is worth being precise about why the earlier fixes could not.
    Un-mixing recovers the true color *given a true alpha*, but the clip arrives 4:2:0, so chroma is shared across
    2x2 blocks and a backdrop pixel next to a dark outline gets its chroma dragged toward neutral. The matte reads
    that as partial coverage. The pixel is really 100% backdrop, so there is no foreground in it to recover — every
    un-mixing gets is a wild extrapolation that clips, leaving the pixel faintly opaque and still magenta.

    The compositor's answer is not to recover that pixel's color but to *replace* it with its neighbours'. An
    alpha-weighted blur samples only foreground — pixels contribute in proportion to how opaque they are — so
    dividing the blurred premultiplied color by the blurred alpha gives the foreground color spread outward across
    the edge. Blend it in by ``1 - alpha`` and the interior is untouched while the band takes the character's own
    color at the character's own edge. Backdrop color then appears nowhere in the output at any alpha.
    """
    if radius <= 0:
        return rgb
    weight = alpha[..., None]
    blurred_color = _box_blur(rgb * weight, radius)
    blurred_weight = _box_blur(weight, radius)
    extended = blurred_color / np.maximum(blurred_weight, 1e-4)
    # Where nothing opaque is in reach, the division is meaningless — keep what is already there rather than
    # smearing a value out of noise.
    reachable = (blurred_weight > 1e-3).astype(np.float32)
    extended = np.clip(extended, 0.0, 1.0) * reachable + rgb * (1.0 - reachable)
    return (rgb * weight + extended * (1.0 - weight)).astype(np.float32)


def _kmeans(points: np.ndarray, k: int, iterations: int = 12, seed: int = 0) -> np.ndarray:
    """Small k-means with k-means++ seeding. Enough for a dozen flat colors out of a few thousand samples."""
    rng = np.random.default_rng(seed)
    centers = [points[rng.integers(len(points))]]
    nearest = ((points - centers[0]) ** 2).sum(1)
    for _ in range(k - 1):
        total = float(nearest.sum())
        if total < 1e-12:
            break
        centers.append(points[rng.choice(len(points), p=nearest / total)])
        nearest = np.minimum(nearest, ((points - centers[-1]) ** 2).sum(1))
    current = np.stack(centers)
    for _ in range(iterations):
        labels = np.argmin(((points[:, None, :] - current[None]) ** 2).sum(2), axis=1)
        for i in range(len(current)):
            member = labels == i
            if member.any():
                current[i] = points[member].mean(0)
    return current.astype(np.float32)


def sample_interior(
    frame: np.ndarray,
    key_rgb: np.ndarray,
    limit: int = 4000,
    seed: int = 0,
) -> np.ndarray:
    """Confidently-interior pixels of one frame, for palette extraction.

    Eroded hard — only pixels whose whole neighbourhood is opaque qualify. The palette must describe the character's
    real flat colors, and a single edge pixel dragged in by 4:2:0 bleed would add a magenta-tinted entry that then
    licenses magenta everywhere.
    """
    alpha = matte(frame, key_rgb)
    interior = _box_blur(alpha[..., None], 3)[..., 0] > 0.995
    points = frame[interior]
    if len(points) == 0:
        return points.reshape(0, 3)
    if len(points) > limit:
        points = points[np.random.default_rng(seed).choice(len(points), limit, replace=False)]
    return points


def extract_palette(
    samples: np.ndarray,
    size: int = 16,
    key_rgb: np.ndarray | None = None,
    min_distance: float = 0.30,
) -> np.ndarray:
    """The character's flat palette, as ``(k, 3)`` float RGB.

    **No entry may sit near the backdrop**, and that guard is what makes `palette_matte` safe rather than
    catastrophic. Coverage is solved as ``a = <observed - key, P - key> / |P - key|^2``; if some ``P`` is close to
    the key that denominator goes to zero, so ordinary backdrop noise divides up to ``a = 1`` and the entire
    background turns *opaque* in the backdrop's own color. Observed live: H3's backdrop is not perfectly uniform,
    and a darker patch of magenta carries *less* chroma, so the chroma matte reads it as character, it gets sampled
    into the palette, and the whole frame goes solid.
    """
    points = samples.astype(np.float32).reshape(-1, 3)
    if key_rgb is not None and len(points):
        points = points[np.linalg.norm(points - key_rgb, axis=1) >= min_distance]
    if len(points) == 0:
        return np.zeros((0, 3), np.float32)
    palette = points if len(points) < size else _kmeans(points, size)
    if key_rgb is not None and len(palette):
        # k-means can still land a centroid near the key by averaging edge pixels; drop those too.
        palette = palette[np.linalg.norm(palette - key_rgb, axis=1) >= min_distance]
    return palette.astype(np.float32)


def quantize_to_palette(
    frame: np.ndarray,
    palette: np.ndarray,
    key_rgb: np.ndarray | None = None,
) -> np.ndarray:
    """Snap every pixel of an image to its nearest palette color. The backdrop joins the palette as an exact entry.

    Applied to the character *still, before it is sent to H3*, this makes the art genuinely flat rather than
    nominally flat: gradients, film-grain speckle and the text-to-image model's own dithering all collapse onto a
    dozen colors. Two things follow. The backdrop becomes one exact value instead of a cloud around one, which is
    the assumption the whole key rests on; and a flatter source gives the video encoder large constant-color regions
    to work with, so there is less for chroma subsampling to smear at the edges.
    """
    entries = palette if key_rgb is None else np.concatenate([palette, key_rgb[None, :]])
    if len(entries) == 0:
        return frame
    best_distance = np.full(frame.shape[:2], np.inf, np.float32)
    best_index = np.zeros(frame.shape[:2], np.int32)
    for i, color in enumerate(entries):
        distance = np.linalg.norm(frame - color, axis=-1)
        closer = distance < best_distance
        best_distance = np.where(closer, distance, best_distance)
        best_index = np.where(closer, i, best_index)
    return entries[best_index].astype(np.float32)


def palette_matte(
    frame: np.ndarray,
    palette: np.ndarray,
    key_rgb: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Explain every pixel as *one palette color composited over the backdrop*. Returns ``(alpha, rgb)``.

    This is the only step here that can *guarantee* the backdrop is gone, and it earns that by exploiting something
    the earlier chroma keying ignored: the source art is flat, so the character is made of a dozen known colors and
    nothing in between. Every in-between color in the decoded clip is therefore an artifact of H.264 — a blend of a
    real color with the backdrop — and is exactly what a composite equation describes.

    For each palette color ``P`` the coverage that best explains the pixel is a projection onto the segment from the
    backdrop to that color::

        a = clamp(<observed - key, P - key> / |P - key|^2, 0, 1)

    and the residual says how well it explains it. Take the palette color with the smallest residual.

    Two properties follow, and they are the point:

    * **A pure-backdrop pixel is explained by ``a = 0`` with zero residual, for every palette entry.** It becomes
      fully transparent no matter what its chroma got smeared into — which is precisely the rim that survived
      choking, un-mixing and edge extension, all of which are identities at ``a == 1``.
    * **The output color is always a palette color**, never a blend, so no amount of backdrop can be left in it.

    It also removes the failure mode of `suppress_residual`: a character that genuinely *is* near the key color has
    that color in its palette, so it matches at ``a = 1`` and is preserved rather than eroded.
    """
    offset = frame - key_rgb
    directions = palette.astype(np.float32) - key_rgb
    # A palette entry sitting on the key would make the projection below divide by ~0 and turn backdrop noise into
    # full coverage. `extract_palette` filters these out; this is the belt to that braces.
    usable = np.linalg.norm(directions, axis=1) >= 0.15
    directions, palette = directions[usable], palette[usable]
    if len(palette) == 0:
        return np.zeros(frame.shape[:2], np.float32), np.zeros_like(frame)
    lengths = np.maximum((directions ** 2).sum(1), 1e-8)

    best_residual = np.full(frame.shape[:2], np.inf, np.float32)
    best_alpha = np.zeros(frame.shape[:2], np.float32)
    best_color = np.zeros_like(frame)

    # Looped over the palette rather than broadcast: a (H, W, k, 3) tensor is hundreds of MB at 768^2 and this Space
    # runs on 2 vCPU / 16 GB.
    for i in range(len(palette)):
        direction = directions[i]
        alpha = np.clip((offset @ direction) / lengths[i], 0.0, 1.0)
        residual = np.linalg.norm(offset - alpha[..., None] * direction, axis=-1)
        better = residual < best_residual
        best_residual = np.where(better, residual, best_residual)
        best_alpha = np.where(better, alpha, best_alpha)
        best_color = np.where(better[..., None], palette[i], best_color)

    return best_alpha.astype(np.float32), best_color.astype(np.float32)


def suppress_residual(
    rgb: np.ndarray,
    alpha: np.ndarray,
    key_rgb: np.ndarray,
    low: float = 0.25,
    high: float = 0.45,
) -> np.ndarray:
    """Drive alpha to zero wherever the *final* color is still the backdrop.

    This is the pass that actually removes the outline, and it exists because `unmix` makes that outline worse
    rather than better. Un-mixing solves ``F = (observed - (1-a)*key)/a``. Feed it a pixel that is 100% backdrop but
    which the matte mis-read as half-covered and it returns the key color at *full* brightness — the pixel goes from
    a muted magenta at 50% opacity to a saturated magenta at 30%, which is more visible, not less. Every alpha-side
    tool is blind to this: choking only moves alpha, and both `unmix` and `extend_edges` are identities at ``a == 1``,
    so a rim pixel the matte called opaque passes through all of them untouched.

    So this works on the axis that is actually wrong: how far a pixel's chroma has gone along the backdrop's own
    direction, as a fraction of the backdrop's chroma.

    **It is gated to the silhouette boundary, and that gate is not optional.** Color alone cannot identify the rim.
    Measured across a magenta-keyed frog and a rose-keyed orange cat, the cat's *own fur* scores higher on this
    fraction (10.9k pixels above 0.7) than the frog's visible magenta rim does (379) — a warm subject genuinely
    shares the +Cr axis with a red-magenta backdrop, so any global threshold that removes the rim also dissolves the
    character. What separates them is position: the rim is backdrop that survived *at the edge of the matte*, while
    the fur is interior. Blurred alpha measures exactly that, so suppression only applies where transparency is
    within a pixel or two, and interior color of any hue is untouched.
    """
    key_chroma = to_ycc(key_rgb[None, None, :])[0, 0][1:]
    magnitude = float(np.linalg.norm(key_chroma))
    if magnitude < 1e-4:
        return alpha
    direction = key_chroma / magnitude
    fraction = np.einsum("...c,c->...", to_ycc(rgb)[..., 1:], direction) / magnitude
    # High where transparency is nearby, ~0 anywhere inside the character. The radius has to reach several pixels
    # inward: measured across a real edge, the rim ran two pixels deep at alpha 0.85 and 0.99, and a tighter gate
    # had already decayed to 0.16 by the second one and left it standing.
    boundary = _smoothstep(0.05, 0.35, 1.0 - _box_blur(alpha[..., None], 6)[..., 0])
    return (alpha * (1.0 - boundary * _smoothstep(low, high, fraction))).astype(np.float32)


def unmix(frame: np.ndarray, alpha: np.ndarray, key_rgb: np.ndarray, strength: float = 1.0) -> np.ndarray:
    """Recover the character's own color from an edge pixel, given the finished matte.

    An edge pixel is a real alpha composite of the character over the backdrop, ``observed = a*F + (1-a)*key``, so
    the character's color comes back out algebraically: ``F = (observed - (1-a)*key) / a``. That removes the fringe
    rather than trading it for a grey rim, and it is an identity at ``a == 1``, so opaque interior pixels are
    provably untouched.

    That last property is why this replaced a chroma-projection despill. Projection subtracts the backdrop's chroma
    direction from *every* pixel, and a magenta key points almost exactly along +Cr — an axis every warm color
    shares — so an early version of this turned an orange cat green.
    """
    if strength <= 0:
        return frame
    safe = np.maximum(alpha, 1e-3)[..., None]
    unmixed = np.clip((frame - (1.0 - alpha)[..., None] * key_rgb) / safe, 0.0, 1.0)
    # Below a few percent coverage the division mostly amplifies noise, and those pixels are invisible anyway —
    # fade the correction out rather than letting it produce confetti.
    trust = np.clip(alpha / 0.15, 0.0, 1.0)[..., None] * strength
    return (frame * (1.0 - trust) + unmixed * trust).astype(np.float32)


def key_frame(
    frame: np.ndarray,
    key_rgb: np.ndarray,
    tolerance: float = 0.10,
    softness: float = 0.09,
    despill: float = 1.0,
    despeckle: int = 2,
    shrink: float = 0.2,
    feather: float = 0.6,
    palette: np.ndarray | None = None,
) -> np.ndarray:
    """Returns RGBA float 0..1.

    The two alphas are deliberately different, and mixing them up is what leaves a pink rim:

    * ``raw`` is the measured coverage, and it is the only alpha for which ``observed = a*F + (1-a)*key`` holds. It
      is therefore what `unmix` must solve against.
    * ``refined`` is the *output* matte — despeckled, choked, feathered. Feeding it back into the un-mixing equation
      asks it to invert a composite that never happened: a choked rim pixel claims 6% coverage where the real
      composite had 25%, the division extrapolates wildly, gets clipped, and the pixel stays the backdrop's color
      while remaining faintly opaque. That is the halo, reintroduced one step after being removed.

    So: recover color against the truth, then hand out the tidied matte.
    """
    if palette is not None and len(palette):
        raw, matched = palette_matte(frame, palette, key_rgb)
        # A light choke still earns its place: H3's backdrop drifts a little across a frame, so a background pixel
        # that differs slightly from the measured key solves to a small non-zero coverage rather than exactly zero.
        refined = refine_matte(raw, despeckle, min(float(shrink), 0.15), feather)
        # Snap only where the composite is partial. A fully opaque pixel is interior and keeps its original color,
        # so soft shading and small details survive; the transition band takes the palette color, which is what
        # makes a backdrop-colored edge impossible. No choke is applied — palette matting derives coverage from the
        # full RGB rather than from chroma alone, so it does not need one.
        keep = _smoothstep(0.90, 1.0, refined)[..., None]
        rgb = matched * (1.0 - keep) + frame * keep
        return np.concatenate([rgb, refined[..., None]], axis=-1).astype(np.float32)

    raw = matte(frame, key_rgb, tolerance, softness)
    refined = refine_matte(raw, despeckle, shrink, feather)
    rgb = unmix(frame, raw, key_rgb, despill)
    if despill > 0:
        # Suppress before extending: the extension is an alpha-weighted blur, so any backdrop-colored pixel still
        # holding alpha would smear its magenta into the neighbours the extension is supposed to be repairing.
        refined = suppress_residual(rgb, refined, key_rgb)
        rgb = extend_edges(rgb, refined)
    return np.concatenate([rgb, refined[..., None]], axis=-1).astype(np.float32)


def key_clip(
    frames: list[np.ndarray],
    key_rgb: np.ndarray | None = None,
    tolerance: float = 0.10,
    softness: float = 0.09,
    despill: float = 1.0,
    despeckle: int = 2,
    shrink: float = 0.2,
    feather: float = 0.6,
    palette: np.ndarray | None = None,
) -> tuple[list[np.ndarray], np.ndarray, list[int]]:
    """Key every frame against the clip's **median** backdrop, and report which frames that backdrop still describes.

    MiniMax-H3 is asked to hold the backdrop flat and unchanged and mostly does — but not always. A run of this
    produced an idle clip whose backdrop slid from magenta to mint green partway through, and both obvious repairs
    are traps:

    * Keying every frame against frame 0 leaves the drifted frames completely un-keyed — a solid green card in the
      middle of the sprite sheet.
    * Keying each frame against *its own* border is worse. On that clip it did remove the green backdrop, and took
      the frog with it: the character is green too, so the drifted key matched the subject and shredded it.

    A drifted frame is not a keying problem, it is a bad frame. So: key everything against the median backdrop,
    which is stable and correct for the frames that behaved, and return the indices of the frames whose own border
    still matches it. There are ~73 frames and a sheet needs 4-8, so discarding the drifted ones costs nothing.

    Returns ``(keyed frames, key color, reliable indices)``.
    """
    measured = np.stack([detect_key_color(f) for f in frames])
    if key_rgb is None:
        key_rgb = np.median(measured, axis=0)
    # A well-behaved clip drifts by a couple of levels; the magenta->green slide was a distance of ~0.8. Anything
    # past this is a different backdrop, not noise.
    reliable = [i for i in range(len(frames)) if float(np.linalg.norm(measured[i] - key_rgb)) < 0.15]
    keyed = [
        key_frame(f, key_rgb, tolerance, softness, despill, despeckle, shrink, feather, palette) for f in frames
    ]
    return keyed, key_rgb, (reliable or list(range(len(frames))))


# ---------------------------------------------------------------------------
# pose extremes
# ---------------------------------------------------------------------------


def _silhouettes(frames: list[np.ndarray], size: int = 64) -> np.ndarray:
    """``(n, size*size)`` float occupancy, cheap enough to compare all-pairs."""
    out = []
    for f in frames:
        a = Image.fromarray((np.clip(f[..., 3], 0, 1) * 255).astype(np.uint8), "L")
        out.append(np.asarray(a.resize((size, size), Image.BILINEAR), np.float32).ravel() / 255.0)
    return np.stack(out)


def pose_extremes(frames: list[np.ndarray], count: int = 8, candidates: list[int] | None = None) -> list[int]:
    """Pick ``count`` frame indices that read as an animator's extremes.

    Two things make a frame an extreme, and this uses both:

    * **It is held.** Motion slows to nothing where a limb reverses direction, so the local minima of frame-to-frame
      silhouette change are the poses the animation actually sits on. In-betweens are the fast frames.
    * **It is distinct.** Evenly spacing samples in *time* oversamples whatever the character does slowly and can
      return eight near-identical frames. Farthest-point sampling in silhouette space instead maximizes the minimum
      distance between chosen poses, which is the property that makes the packed sheet read as an animation.

    Frame 0 seeds the selection: it is the conditioning image the clip starts and ends on, so it is the rest pose the
    loop returns to and belongs on every sheet.
    """
    n = len(frames)
    # `candidates` excludes frames the keyer could not trust (see `key_clip`). Motion is still measured across the
    # whole clip so speed stays meaningful; only the *choice* is restricted.
    allowed = sorted(set(candidates)) if candidates else list(range(n))
    if len(allowed) <= count:
        return allowed

    sil = _silhouettes(frames)
    motion = np.concatenate([[0.0], np.linalg.norm(np.diff(sil, axis=0), axis=1)])
    travelled = np.cumsum(motion)
    if travelled[-1] <= 0:
        return sorted(set(np.linspace(0, n - 1, count).round().astype(int).tolist()) & set(allowed)) or allowed[:count]

    # Equal steps of *accumulated* silhouette change, so samples land densely where the pose is actually changing
    # and sparsely through the held parts. Sampling uniformly in time instead looks mushy, and preferring the
    # slowest frames — which an earlier version here did — oversamples the idle bounce at the head of the clip.
    allowed_array = np.array(allowed)
    picks = {
        int(allowed_array[np.argmin(np.abs(travelled[allowed_array] - target))])
        for target in np.linspace(0, travelled[-1], count)
    }
    return sorted(picks)


def cycle_extremes(
    frames: list[np.ndarray],
    count: int = 8,
    candidates: list[int] | None = None,
    low: float = 0.2,
    high: float = 0.85,
) -> list[int]:
    """Pick ``count`` indices spanning one *seamless cycle*, for moves that repeat (walk, run).

    A whole-clip selection is wrong for a cycle. The clip starts at the rest pose, accelerates into the gait, and
    settles back at the end, because that is what the first/last-frame pin asks for. Sample across all of that and
    the sheet contains the start-up and the settle, so the walk visibly resets instead of looping.

    Instead: search the sustained middle of the clip for the pair of frames whose silhouettes match most closely —
    that pair *is* the seam of one gait cycle — and sample between them.
    """
    n = len(frames)
    allowed = set(candidates) if candidates else set(range(n))
    window = [i for i in range(int(n * low), int(n * high)) if i in allowed]
    if len(window) <= count:
        return pose_extremes(frames, count, candidates)

    sil = _silhouettes(frames)
    min_gap = max(6, n // 10)
    max_gap = max(min_gap + 1, n // 2)

    best = None
    for i in window:
        for j in window:
            gap = j - i
            if not (min_gap <= gap <= max_gap):
                continue
            distance = float(np.abs(sil[i] - sil[j]).mean())
            if best is None or distance < best[0]:
                best = (distance, i, j)
    if best is None:
        return pose_extremes(frames, count, candidates)

    _distance, start, end = best
    return sorted({start + round(k * (end - start) / count) for k in range(count)})


# ---------------------------------------------------------------------------
# trimming and packing
# ---------------------------------------------------------------------------


def content_bbox(frames: list[np.ndarray], threshold: float = 0.15, pad: int = 2) -> tuple[int, int, int, int]:
    """The union of every frame's opaque area.

    One box for the whole move, never per frame: cropping each pose to its own bounds re-centers the character on
    every frame and the sprite jitters when played back. The union keeps the character's position within the cell
    exactly as generated, so a walk cycle stays planted.
    """
    h, w = frames[0].shape[:2]
    mask = np.zeros((h, w), bool)
    for f in frames:
        mask |= f[..., 3] > threshold
    if not mask.any():
        return 0, 0, w, h
    rows, cols = np.where(mask.any(1))[0], np.where(mask.any(0))[0]
    return (
        max(0, int(cols[0]) - pad),
        max(0, int(rows[0]) - pad),
        min(w, int(cols[-1]) + 1 + pad),
        min(h, int(rows[-1]) + 1 + pad),
    )


def to_pil(rgba: np.ndarray) -> Image.Image:
    return Image.fromarray((np.clip(rgba, 0, 1) * 255).astype(np.uint8), "RGBA")


def fit_height(sprites: list[Image.Image], max_height: int) -> list[Image.Image]:
    """Scale a move down so its cell is at most ``max_height`` tall, or return it untouched.

    One ratio for the whole move, derived from the shared crop box, so every frame keeps the same scale and the
    character does not appear to grow between poses. Only ever downscales — upscaling a sprite past what the clip
    actually resolved just bakes in interpolation.
    """
    if not sprites or max_height <= 0:
        return sprites
    height = sprites[0].size[1]
    if height <= max_height:
        return sprites
    ratio = max_height / height
    size = (max(1, round(sprites[0].size[0] * ratio)), max_height)
    return [s.resize(size, Image.LANCZOS) for s in sprites]


@dataclass
class Move:
    """One animation: a named clip already cut down to its extremes."""

    name: str
    frames: list[Image.Image] = field(default_factory=list)
    fps: float = 12.0

    @property
    def size(self) -> tuple[int, int]:
        return self.frames[0].size if self.frames else (0, 0)


def pack_sheet(moves: list[Move], columns: int | None = None, margin: int = 0) -> tuple[Image.Image, dict]:
    """Pack every move onto one uniform grid, one row per move.

    A uniform cell — the largest frame across every move, not per move — is what lets a runtime address a frame as
    ``(row, column)`` arithmetic with no per-frame lookup. Sprites are bottom-center anchored inside their cell so a
    character whose kick is wider than its idle still stands on the same ground line.
    """
    if not moves:
        raise ValueError("No moves to pack.")

    cell_w = max(m.size[0] for m in moves) + margin * 2
    cell_h = max(m.size[1] for m in moves) + margin * 2
    cols = columns or max(len(m.frames) for m in moves)
    rows = sum(math.ceil(len(m.frames) / cols) for m in moves)

    sheet = Image.new("RGBA", (cell_w * cols, cell_h * rows), (0, 0, 0, 0))
    atlas = {
        "cell": {"width": cell_w, "height": cell_h},
        "columns": cols,
        "anchor": "bottom-center",
        "moves": {},
    }

    row = 0
    for move in moves:
        entries = []
        for i, frame in enumerate(move.frames):
            c, r = i % cols, row + i // cols
            x = c * cell_w + (cell_w - frame.size[0]) // 2
            y = r * cell_h + (cell_h - frame.size[1])  # bottom aligned: the ground line is shared
            sheet.paste(frame, (x, y))
            entries.append({"index": c + (r - row) * cols, "row": r, "col": c})
        atlas["moves"][move.name] = {
            "row": row,
            "frames": len(move.frames),
            "fps": move.fps,
            "cells": entries,
        }
        row += math.ceil(len(move.frames) / cols)

    atlas["sheet"] = {"width": sheet.size[0], "height": sheet.size[1], "rows": rows}
    return sheet, atlas


# ---------------------------------------------------------------------------
# encoding
# ---------------------------------------------------------------------------


def save_animation(frames: list[Image.Image], path: str, fps: float = 12.0) -> str:
    """Write a looping animation, preserving alpha.

    WebP, because an animated GIF has a 1-bit mask: every partially transparent edge pixel becomes a hard on/off
    decision and the sprite gets a visible fringe of whatever backdrop the encoder guessed.
    """
    duration = max(20, int(round(1000.0 / max(fps, 1e-3))))
    frames[0].save(
        path,
        save_all=True,
        append_images=frames[1:],
        duration=duration,
        loop=0,
        lossless=True,
        disposal=2,
    )
    return path


def read_video(path: str, max_frames: int = 240) -> tuple[list[np.ndarray], float]:
    """Decode a clip to float RGB frames. Uses imageio-ffmpeg, which ships its own ffmpeg binary."""
    import imageio.v3 as iio

    meta = iio.immeta(path, plugin="pyav")
    fps = float(meta.get("fps") or 24.0)
    frames = []
    for i, frame in enumerate(iio.imiter(path, plugin="pyav")):
        if i >= max_frames:
            break
        frames.append(np.asarray(frame, np.float32)[..., :3] / 255.0)
    return frames, fps


def data_uri(image: Image.Image) -> str:
    import io

    buf = io.BytesIO()
    image.save(buf, "PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()


def write_atlas(atlas: dict, path: str) -> str:
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(atlas, handle, indent=2)
    return path