studio/rigging/deform.py

"""studio.rigging.deform — forward kinematics + Linear Blend Skinning warp.

Pose representation: a dict mapping joint_name → local rotation (radians). Local
rotations propagate down the tree to produce a world rotation per joint; each
joint's deformed position is its parent's deformed position plus the rest offset
rotated by the parent's world rotation.

Pixel warping is LBS forward-splat: each foreground pixel computes its blended
output position as Σ_b weight_b · T_b(pixel), then bilinearly splats its color
into the output buffer. Subject preservation is structural: every output pixel
is a weighted combination of source pixels — no pixel is invented.
"""
from __future__ import annotations

from typing import Mapping, Tuple

import numpy as np

from pixel_cursor.rigging import BONES, JOINT_INDEX, NUM_JOINTS, PARENTS, Skeleton


def _eval_order() -> Tuple[int, ...]:
    root_idx = PARENTS.index(-1)
    order = [root_idx]
    queue = [root_idx]
    seen = {root_idx}
    while queue:
        parent = queue.pop(0)
        for child, parent_idx in enumerate(PARENTS):
            if parent_idx == parent and child not in seen:
                order.append(child)
                seen.add(child)
                queue.append(child)
    return tuple(order)


_EVAL_ORDER: Tuple[int, ...] = _eval_order()


def forward_kinematics(
    rest: Skeleton,
    local_rotations: Mapping[str, float] | None = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """Run FK from a dict of local rotations.

    Returns (deformed_positions shape (14, 2), world_rotations shape (14,)).
    """
    local = np.zeros(NUM_JOINTS, dtype=np.float32)
    if local_rotations:
        for name, angle in local_rotations.items():
            local[JOINT_INDEX[name]] = float(angle)

    deformed = rest.positions.copy().astype(np.float32)
    world_rot = np.zeros(NUM_JOINTS, dtype=np.float32)

    root_idx = PARENTS.index(-1)
    world_rot[root_idx] = local[root_idx]

    for joint in _EVAL_ORDER:
        if joint == root_idx:
            continue
        parent = PARENTS[joint]
        world_rot[joint] = world_rot[parent] + local[joint]
        rest_offset = rest.positions[joint] - rest.positions[parent]
        c, s = np.cos(world_rot[parent]), np.sin(world_rot[parent])
        rotated_y = c * rest_offset[0] - s * rest_offset[1]
        rotated_x = s * rest_offset[0] + c * rest_offset[1]
        deformed[joint, 0] = deformed[parent, 0] + rotated_y
        deformed[joint, 1] = deformed[parent, 1] + rotated_x

    return deformed.astype(np.float32), world_rot


def deform_sprite(
    image: np.ndarray,
    mask: np.ndarray,
    rest: Skeleton,
    weights: np.ndarray,
    *,
    local_rotations: Mapping[str, float] | None = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """Apply LBS warp to a sprite. Returns (deformed_rgba (H, W, 4), deformed_mask (H, W) bool).

    Algorithm:
      1. FK produces deformed joint positions + per-joint world rotations.
      2. For each foreground source pixel p:
            T_b(p) = deformed[parent(b)] + R(world_rot[parent(b)]) · (p - rest[parent(b)])
            out_pos = Σ_b weights[p, b] · T_b(p)
      3. Bilinearly splat source RGBA into the output accumulator at out_pos.
      4. Normalise accumulator by accumulated weight per output pixel.
    """
    if mask.ndim != 2:
        raise ValueError(f"mask must be 2D, got {mask.shape}")
    H, W = mask.shape

    deformed_pos, world_rot = forward_kinematics(rest, local_rotations)

    if image.ndim == 2:
        rgba_in = np.stack(
            [image, image, image, np.full_like(image, 255)], axis=-1
        )
    elif image.shape[-1] == 3:
        rgba_in = np.concatenate(
            [image, np.full(image.shape[:2] + (1,), 255, dtype=np.uint8)], axis=-1
        )
    elif image.shape[-1] == 4:
        rgba_in = image
    else:
        raise ValueError(f"unsupported image shape {image.shape}")

    mask_bool = mask.astype(bool)
    ys, xs = np.where(mask_bool)
    if ys.size == 0:
        return np.zeros((H, W, 4), dtype=np.uint8), np.zeros((H, W), dtype=bool)

    src_coords = np.stack([ys, xs], axis=-1).astype(np.float32)
    src_colors = rgba_in[ys, xs].astype(np.float32)
    src_weights = weights[ys, xs]  # (N_fg, N_BONES)

    out_pos = np.zeros_like(src_coords)
    for b, (parent_idx, _child_idx) in enumerate(BONES):
        rot = float(world_rot[parent_idx])
        c, s = np.cos(rot), np.sin(rot)
        offset = src_coords - rest.positions[parent_idx]
        ry = c * offset[:, 0] - s * offset[:, 1]
        rx = s * offset[:, 0] + c * offset[:, 1]
        rotated = np.stack([ry, rx], axis=-1)
        per_bone_out = deformed_pos[parent_idx] + rotated
        out_pos += src_weights[:, b : b + 1] * per_bone_out

    accum = np.zeros((H, W, 4), dtype=np.float32)
    wsum = np.zeros((H, W), dtype=np.float32)

    oy = out_pos[:, 0]
    ox = out_pos[:, 1]
    oy_int = np.floor(oy).astype(np.int32)
    ox_int = np.floor(ox).astype(np.int32)
    fy = oy - oy_int
    fx = ox - ox_int

    for dy in (0, 1):
        for dx in (0, 1):
            wy = (1.0 - fy) if dy == 0 else fy
            wx = (1.0 - fx) if dx == 0 else fx
            wgt = wy * wx
            yy = oy_int + dy
            xx = ox_int + dx
            valid = (yy >= 0) & (yy < H) & (xx >= 0) & (xx < W) & (wgt > 0)
            if not valid.any():
                continue
            yy_v = yy[valid]
            xx_v = xx[valid]
            w_v = wgt[valid]
            for ch in range(4):
                np.add.at(accum[..., ch], (yy_v, xx_v), w_v * src_colors[valid, ch])
            np.add.at(wsum, (yy_v, xx_v), w_v)

    deformed_rgba = np.zeros((H, W, 4), dtype=np.uint8)
    has_color = wsum > 0
    np.divide(
        accum,
        np.maximum(wsum[..., None], 1e-12),
        out=accum,
        where=has_color[..., None],
    )
    deformed_rgba[has_color] = np.clip(accum[has_color], 0, 255).astype(np.uint8)

    return deformed_rgba, has_color


__all__ = ["forward_kinematics", "deform_sprite"]