studio/rigging/parts.py

"""studio.rigging.parts — piece-mask sidecar (Phase 3 / Session 3.5).

Each sprite gets two artifacts alongside its joints sidecar:

    <sprite>.png            (the sprite image)
    <sprite>.joints.json    (joint annotations, Session 2)
    <sprite>.parts.png      (uint8 grayscale; pixel = piece_id; 0 = background)
    <sprite>.parts.json     (piece index, bone binding, z-order, kind)

The parts.png is mode='L' so every editor opens and round-trips it cleanly. A
separate <sprite>.parts.preview.png with a color palette is generated for human
inspection; it is NOT load-bearing.

JSON format (versioned, forwards-compatible):

    {
        "schema_version": 1,
        "image": "goblin.png",
        "image_shape": [256, 256],
        "pieces": {
            "1":  {"name": "torso",       "bone": "neck",       "z": 1, "kind": "body"},
            ...
            "14": {"name": "sword",       "bone": "r_wrist",    "z": 9, "kind": "attachment"}
        }
    }

Piece IDs are 1-indexed. ID 0 is reserved for background (no piece).

`bone` is the joint name the piece pivots around. Under forward kinematics, the
piece rotates rigidly around that joint's position by that joint's world
rotation. `z` is the render order — higher draws on top. `kind` is "body" or
"attachment"; the default rigid-piece engine renders body-only.
"""
from __future__ import annotations

import colorsys
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional, Tuple

import numpy as np
from PIL import Image

from pixel_cursor.rigging import JOINT_INDEX, JOINT_NAMES


SCHEMA_VERSION: int = 1


CANONICAL_HUMANOID_PIECES: Tuple[Tuple[str, str, int, str], ...] = (
    # (name, bone (joint name), z-order, kind)
    ("torso",       "neck",       1, "body"),
    ("l_thigh",     "l_hip",      2, "body"),
    ("r_thigh",     "r_hip",      2, "body"),
    ("l_shin",      "l_knee",     3, "body"),
    ("r_shin",      "r_knee",     3, "body"),
    ("l_foot",      "l_ankle",    4, "body"),
    ("r_foot",      "r_ankle",    4, "body"),
    ("l_upper_arm", "l_shoulder", 5, "body"),
    ("r_upper_arm", "r_shoulder", 5, "body"),
    ("l_forearm",   "l_elbow",    6, "body"),
    ("r_forearm",   "r_elbow",    6, "body"),
    ("l_hand",      "l_wrist",    7, "body"),
    ("r_hand",      "r_wrist",    7, "body"),
    ("head",        "head",       8, "body"),
)


@dataclass(frozen=True)
class Piece:
    piece_id: int
    name: str
    bone: str
    z: int
    kind: str


@dataclass(frozen=True)
class PartsMask:
    id_map: np.ndarray
    pieces: Tuple[Piece, ...]
    image_shape: Tuple[int, int]

    def piece_by_name(self, name: str) -> Optional[Piece]:
        for p in self.pieces:
            if p.name == name:
                return p
        return None

    def piece_by_id(self, piece_id: int) -> Optional[Piece]:
        for p in self.pieces:
            if p.piece_id == piece_id:
                return p
        return None

    @property
    def body_pieces(self) -> Tuple[Piece, ...]:
        return tuple(p for p in self.pieces if p.kind == "body")

    @property
    def attachment_pieces(self) -> Tuple[Piece, ...]:
        return tuple(p for p in self.pieces if p.kind == "attachment")


def parts_png_path_for(image_path: Path | str) -> Path:
    p = Path(image_path)
    return p.with_name(p.stem + ".parts.png")


def parts_json_path_for(image_path: Path | str) -> Path:
    p = Path(image_path)
    return p.with_name(p.stem + ".parts.json")


def parts_preview_path_for(image_path: Path | str) -> Path:
    p = Path(image_path)
    return p.with_name(p.stem + ".parts.preview.png")


def load_parts(image_path: Path | str) -> Optional[PartsMask]:
    """Return the PartsMask for image_path, or None if either sidecar is missing.

    Raises ValueError if the JSON exists but is malformed or has a higher
    schema_version than we support.
    """
    png_path = parts_png_path_for(image_path)
    json_path = parts_json_path_for(image_path)
    if not png_path.exists() or not json_path.exists():
        return None

    data = json.loads(json_path.read_text())
    version = int(data.get("schema_version", 1))
    if version > SCHEMA_VERSION:
        raise ValueError(
            f"parts sidecar schema_version {version} > supported {SCHEMA_VERSION} "
            f"({json_path})"
        )

    shape_field = data["image_shape"]
    H, W = int(shape_field[0]), int(shape_field[1])

    img = Image.open(png_path)
    if img.mode != "L":
        img = img.convert("L")
    arr = np.asarray(img, dtype=np.uint8)
    if arr.shape != (H, W):
        raise ValueError(
            f"parts.png shape {arr.shape} != image_shape {(H, W)} ({png_path})"
        )

    pieces_field = data.get("pieces", {})
    pieces = []
    for pid_str, entry in pieces_field.items():
        pid = int(pid_str)
        if pid == 0:
            raise ValueError(
                f"piece_id 0 is reserved for background ({json_path})"
            )
        if pid < 0 or pid > 255:
            raise ValueError(
                f"piece_id {pid} out of uint8 range ({json_path})"
            )
        bone = entry["bone"]
        if bone not in JOINT_INDEX:
            raise ValueError(
                f"piece '{entry.get('name', pid)}' bone '{bone}' is not a known "
                f"joint ({json_path}); known: {list(JOINT_NAMES)}"
            )
        pieces.append(Piece(
            piece_id=pid,
            name=str(entry["name"]),
            bone=bone,
            z=int(entry.get("z", 0)),
            kind=str(entry.get("kind", "body")),
        ))

    return PartsMask(
        id_map=arr,
        pieces=tuple(pieces),
        image_shape=(H, W),
    )


def save_parts(
    image_path: Path | str,
    id_map: np.ndarray,
    pieces: Tuple[Piece, ...],
    *,
    image_name: Optional[str] = None,
) -> Tuple[Path, Path]:
    """Write <sprite>.parts.png + <sprite>.parts.json. Returns (png_path, json_path)."""
    if id_map.ndim != 2:
        raise ValueError(f"id_map must be 2D, got {id_map.shape}")
    if id_map.dtype != np.uint8:
        raise ValueError(f"id_map must be uint8, got {id_map.dtype}")
    png_path = parts_png_path_for(image_path)
    json_path = parts_json_path_for(image_path)
    name = image_name if image_name is not None else Path(image_path).name

    Image.fromarray(id_map, mode="L").save(png_path)

    payload = {
        "schema_version": SCHEMA_VERSION,
        "image": name,
        "image_shape": [int(id_map.shape[0]), int(id_map.shape[1])],
        "pieces": {
            str(p.piece_id): {
                "name": p.name,
                "bone": p.bone,
                "z": p.z,
                "kind": p.kind,
            }
            for p in pieces
        },
    }
    json_path.write_text(json.dumps(payload, indent=2, sort_keys=False) + "\n")
    return png_path, json_path


def save_preview(
    image_path: Path | str,
    parts: PartsMask,
) -> Path:
    """Write a color-coded RGBA preview alongside the parts sidecar.

    Pure operator convenience — not load-bearing. Background is transparent.
    """
    H, W = parts.image_shape
    rgba = np.zeros((H, W, 4), dtype=np.uint8)
    palette = _build_preview_palette(parts.pieces)
    for piece in parts.pieces:
        mask = parts.id_map == piece.piece_id
        if not mask.any():
            continue
        rgba[mask] = palette[piece.piece_id]
    out_path = parts_preview_path_for(image_path)
    Image.fromarray(rgba, mode="RGBA").save(out_path)
    return out_path


def _build_preview_palette(
    pieces: Tuple[Piece, ...],
) -> Dict[int, Tuple[int, int, int, int]]:
    palette: Dict[int, Tuple[int, int, int, int]] = {}
    n = max(1, len(pieces))
    for i, p in enumerate(pieces):
        hue = (i / n) % 1.0
        sat = 0.65
        val = 0.95 if p.kind == "body" else 0.55
        r, g, b = colorsys.hsv_to_rgb(hue, sat, val)
        palette[p.piece_id] = (int(r * 255), int(g * 255), int(b * 255), 220)
    return palette


def canonical_humanoid_pieces(
    extras: Tuple[Tuple[str, str, int, str], ...] = (),
) -> Tuple[Piece, ...]:
    """Return the canonical 13 humanoid pieces plus any extras (attachments).

    extras: iterable of (name, bone, z, kind) tuples to append. Piece IDs are
    assigned sequentially starting from 1.
    """
    pieces = []
    next_id = 1
    for name, bone, z, kind in CANONICAL_HUMANOID_PIECES:
        pieces.append(Piece(piece_id=next_id, name=name, bone=bone, z=z, kind=kind))
        next_id += 1
    for name, bone, z, kind in extras:
        pieces.append(Piece(piece_id=next_id, name=name, bone=bone, z=z, kind=kind))
        next_id += 1
    return tuple(pieces)


__all__ = [
    "SCHEMA_VERSION",
    "CANONICAL_HUMANOID_PIECES",
    "Piece",
    "PartsMask",
    "parts_png_path_for",
    "parts_json_path_for",
    "parts_preview_path_for",
    "load_parts",
    "save_parts",
    "save_preview",
    "canonical_humanoid_pieces",
]