snap2sim/schema.py

"""Shared data contract between vision analysis and scene generation."""

from __future__ import annotations

import math
from typing import Any, Literal

MotionType = Literal["rotate", "translate", "oscillate", "static", "screw", "orbit", "pulse"]


EXAMPLE_ANALYSIS: dict[str, Any] = {
    "component": "ratchet wrench head",
    "confidence": 0.72,
    "summary": (
        "A ratchet head converts a back-and-forth handle motion into one-way "
        "socket rotation. A pawl locks against the gear teeth during the drive "
        "stroke and slips over the teeth during the return stroke."
    ),
    "trigger": "handle swings clockwise and counterclockwise",
    "motion_sequence": [
        "handle applies torque to the outer head",
        "spring-loaded pawl locks into the ratchet gear",
        "gear and socket rotate on the drive stroke",
        "pawl rides over gear teeth on the return stroke",
    ],
    "parts": [
        {
            "id": "housing",
            "name": "head housing",
            "role": "keeps the mechanism aligned",
            "geometry": {
                "shape": "cylinder",
                "size": [3.2, 0.55, 3.2],
                "position": [0, 0, 0],
                "rotation": [1.5708, 0, 0],
                "color": "steel",
            },
            "motion": {"type": "static"},
            "annotation": {
                "point": [0.5, 0.55],
                "label": "head housing",
                "note": "visible outer shell around the ratchet mechanism",
            },
        },
        {
            "id": "ratchet_gear",
            "name": "ratchet gear",
            "role": "carries the socket and teeth",
            "geometry": {
                "shape": "gear",
                "teeth": 24,
                "size": [1.45, 0.4, 1.45],
                "position": [0, 0.15, 0],
                "color": "cyan",
            },
            "motion": {
                "type": "rotate",
                "axis": [0, 1, 0],
                "speed": 0.95,
                "phase": 0,
            },
            "annotation": {
                "point": [0.5, 0.48],
                "label": "ratchet gear",
                "note": "central rotating socket carrier",
            },
        },
        {
            "id": "pawl",
            "name": "spring pawl",
            "role": "locks and releases against gear teeth",
            "geometry": {
                "shape": "cone",
                "size": [0.34, 0.62, 0.34],
                "position": [1.25, 0.3, 0.18],
                "rotation": [0.55, 0, -0.75],
                "color": "amber",
            },
            "motion": {
                "type": "oscillate",
                "axis": [0, 1, 0],
                "amplitude": 0.28,
                "speed": 3.2,
                "phase": 0.6,
            },
            "annotation": {
                "point": [0.67, 0.44],
                "label": "spring pawl",
                "note": "locking tooth that alternates grip and slip",
            },
        },
        {
            "id": "pawl_spring",
            "name": "pawl return spring",
            "role": "pushes the pawl back into the gear teeth",
            "geometry": {
                "shape": "spring",
                "coils": 6,
                "wire": 0.035,
                "size": [0.34, 0.9, 0.34],
                "position": [1.55, 0.32, -0.2],
                "rotation": [0.2, 0, 0.45],
                "color": "orange",
            },
            "motion": {"type": "pulse", "speed": 2.4, "amplitude": 0.08},
            "annotation": {
                "point": [0.72, 0.5],
                "label": "return spring",
                "note": "small spring biasing the locking pawl",
            },
        },
        {
            "id": "selector_pin",
            "name": "selector pin",
            "role": "slides the selector and changes pawl bias direction",
            "geometry": {
                "shape": "capsule",
                "size": [0.22, 0.85, 0.22],
                "position": [-0.95, 0.56, -0.95],
                "rotation": [1.5708, 0, 0.25],
                "color": "steel",
            },
            "motion": {
                "type": "translate",
                "axis": [1, 0, 0],
                "range": [-0.08, 0.08],
                "speed": 1.3,
            },
            "annotation": {
                "point": [0.38, 0.3],
                "label": "selector pin",
                "note": "rounded pin under the direction selector",
            },
        },
        {
            "id": "retaining_ring",
            "name": "retaining ring",
            "role": "keeps the socket gear captured in the head",
            "geometry": {
                "shape": "torus",
                "size": [1.9, 0.18, 1.9],
                "position": [0, 0.43, 0],
                "color": "cyan",
            },
            "motion": {"type": "static"},
            "annotation": {
                "point": [0.5, 0.36],
                "label": "retaining ring",
                "note": "circular clip around the socket drive",
            },
        },
    ],
}


ANALYSIS_SCHEMA: dict[str, Any] = {
    "type": "object",
    "required": ["component", "summary", "trigger", "motion_sequence", "parts"],
    "properties": {
        "component": {"type": "string"},
        "confidence": {"type": "number", "minimum": 0, "maximum": 1},
        "render_mode": {"type": "string", "enum": ["three", "annotate", "unavailable"]},
        "summary": {"type": "string"},
        "trigger": {"type": "string"},
        "motion_sequence": {"type": "array", "items": {"type": "string"}},
        "parts": {
            "type": "array",
            "maxItems": 6,
            "items": {
                "type": "object",
                "required": ["id", "name", "role"],
                "properties": {
                    "id": {"type": "string"},
                    "name": {"type": "string"},
                    "role": {"type": "string"},
                    "geometry": {
                        "type": "object",
                        "required": ["shape", "size", "position"],
                        "properties": {
                            "shape": {
                                "type": "string",
                                "enum": [
                                    "box",
                                    "cylinder",
                                    "sphere",
                                    "gear",
                                    "rod",
                                    "cone",
                                    "capsule",
                                    "torus",
                                    "spring",
                                ],
                            },
                            "size": {
                                "type": "array",
                                "items": {"type": "number"},
                                "minItems": 3,
                                "maxItems": 3,
                            },
                            "position": {
                                "type": "array",
                                "items": {"type": "number"},
                                "minItems": 3,
                                "maxItems": 3,
                            },
                            "rotation": {
                                "type": "array",
                                "items": {"type": "number"},
                                "minItems": 3,
                                "maxItems": 3,
                            },
                            "teeth": {"type": "integer", "minimum": 6, "maximum": 80},
                            "coils": {"type": "integer", "minimum": 2, "maximum": 12},
                            "wire": {"type": "number"},
                            "color": {"type": "string"},
                        },
                    },
                    "motion": {
                        "type": "object",
                        "required": ["type"],
                        "properties": {
                            "type": {
                                "type": "string",
                                "enum": [
                                    "rotate",
                                    "translate",
                                    "oscillate",
                                    "static",
                                    "screw",
                                    "orbit",
                                    "pulse",
                                ],
                            },
                            "axis": {
                                "type": "array",
                                "items": {"type": "number"},
                                "minItems": 3,
                                "maxItems": 3,
                            },
                            "speed": {"type": "number"},
                            "amplitude": {"type": "number"},
                            "phase": {"type": "number"},
                            "range": {
                                "type": "array",
                                "items": {"type": "number"},
                                "minItems": 2,
                                "maxItems": 2,
                            },
                            "pitch": {"type": "number"},
                            "pivot": {
                                "type": "array",
                                "items": {"type": "number"},
                                "minItems": 3,
                                "maxItems": 3,
                            },
                        },
                    },
                    "annotation": {
                        "type": "object",
                        "required": ["point"],
                        "properties": {
                            "point": {
                                "type": "array",
                                "items": {"type": "number", "minimum": 0, "maximum": 1},
                                "minItems": 2,
                                "maxItems": 2,
                            },
                            "box": {
                                "type": "array",
                                "items": {"type": "number", "minimum": 0, "maximum": 1},
                                "minItems": 4,
                                "maxItems": 4,
                            },
                            "label": {"type": "string"},
                            "note": {"type": "string"},
                        },
                    },
                },
            },
        },
    },
}


_SHAPES = {"box", "cylinder", "sphere", "gear", "rod", "cone", "capsule", "torus", "spring"}
_MOTIONS = {"rotate", "translate", "oscillate", "static", "screw", "orbit", "pulse"}
_RENDER_MODES = {"three", "annotate", "unavailable"}
DEFAULT_CONFIDENCE_THRESHOLD = 0.5


def normalize_confidence_threshold(threshold: Any = DEFAULT_CONFIDENCE_THRESHOLD) -> float:
    if isinstance(threshold, bool) or threshold is None:
        return DEFAULT_CONFIDENCE_THRESHOLD
    try:
        value = float(threshold)
    except (TypeError, ValueError):
        return DEFAULT_CONFIDENCE_THRESHOLD
    if not math.isfinite(value):
        return DEFAULT_CONFIDENCE_THRESHOLD
    return max(0.0, min(1.0, value))


def select_render_mode(
    analysis: dict[str, Any],
    threshold: Any = DEFAULT_CONFIDENCE_THRESHOLD,
) -> str:
    """Pick the safest advisory renderer for a validated analysis payload."""
    threshold = normalize_confidence_threshold(threshold)
    explicit_mode = analysis.get("render_mode")

    parts = analysis.get("parts") if isinstance(analysis, dict) else []
    has_geometry = any(isinstance(part, dict) and isinstance(part.get("geometry"), dict) for part in parts)
    has_annotation = any(isinstance(part, dict) and isinstance(part.get("annotation"), dict) for part in parts)
    confidence = analysis.get("confidence", 1)
    low_confidence = _is_number(confidence) and float(confidence) < threshold
    if has_annotation and (low_confidence or not has_geometry):
        mode = "annotate"
    elif has_geometry and not low_confidence:
        mode = "three"
    elif has_annotation:
        mode = "annotate"
    else:
        mode = "unavailable"
    if explicit_mode in _RENDER_MODES:
        return _lower_render_mode(str(explicit_mode), mode)
    return mode


def _lower_render_mode(left: str, right: str) -> str:
    rank = {"unavailable": 0, "annotate": 1, "three": 2}
    return left if rank[left] <= rank[right] else right


def validate_analysis(payload: dict[str, Any]) -> dict[str, Any]:
    """Validate a model analysis payload against the scene-generation contract."""
    if not isinstance(payload, dict):
        raise ValueError("Analysis payload must be a JSON object.")

    for key in ["component", "summary", "trigger"]:
        _require_string(payload, key, key)
    _require_string_list(payload, "motion_sequence", "motion_sequence")
    if "render_mode" in payload and payload["render_mode"] not in _RENDER_MODES:
        raise ValueError("Invalid analysis payload at render_mode: unsupported mode")

    parts = payload.get("parts")
    if not isinstance(parts, list) or not parts:
        raise ValueError("Invalid analysis payload at parts: expected a non-empty list")
    if len(parts) > 6:
        raise ValueError("Invalid analysis payload at parts: expected no more than 6 items")

    for index, part in enumerate(parts):
        path = f"parts.{index}"
        if not isinstance(part, dict):
            raise ValueError(f"Invalid analysis payload at {path}: expected an object")
        for key in ["id", "name", "role"]:
            _require_string(part, key, f"{path}.{key}")

        has_geometry = False
        geometry = part.get("geometry")
        if geometry is not None:
            if not isinstance(geometry, dict):
                raise ValueError(f"Invalid analysis payload at {path}.geometry: expected an object")
            shape = geometry.get("shape")
            if shape not in _SHAPES:
                raise ValueError(f"Invalid analysis payload at {path}.geometry.shape: unsupported shape")
            _require_number_list(geometry, "size", f"{path}.geometry.size", 3)
            _require_number_list(geometry, "position", f"{path}.geometry.position", 3)
            if "rotation" in geometry:
                _require_number_list(geometry, "rotation", f"{path}.geometry.rotation", 3)
            if "teeth" in geometry and not isinstance(geometry["teeth"], int):
                raise ValueError(f"Invalid analysis payload at {path}.geometry.teeth: expected an integer")
            if "coils" in geometry and not isinstance(geometry["coils"], int):
                raise ValueError(f"Invalid analysis payload at {path}.geometry.coils: expected an integer")
            if "wire" in geometry and not _is_number(geometry["wire"]):
                raise ValueError(f"Invalid analysis payload at {path}.geometry.wire: expected a number")
            has_geometry = True

        has_annotation = False
        annotation = part.get("annotation")
        if annotation is not None:
            if not isinstance(annotation, dict):
                raise ValueError(f"Invalid analysis payload at {path}.annotation: expected an object")
            _require_unit_number_list(annotation, "point", f"{path}.annotation.point", 2)
            if "box" in annotation:
                _require_unit_number_list(annotation, "box", f"{path}.annotation.box", 4)
            for key in ["label", "note"]:
                if key in annotation and not isinstance(annotation[key], str):
                    raise ValueError(f"Invalid analysis payload at {path}.annotation.{key}: expected a string")
            has_annotation = True

        if not has_geometry and not has_annotation:
            raise ValueError(
                f"Invalid analysis payload at {path}: expected geometry or annotation"
            )

        motion = part.get("motion")
        if motion is not None:
            if not isinstance(motion, dict):
                raise ValueError(f"Invalid analysis payload at {path}.motion: expected an object")
            motion_type = motion.get("type")
            if motion_type not in _MOTIONS:
                raise ValueError(f"Invalid analysis payload at {path}.motion.type: unsupported motion")
            if "axis" in motion:
                _require_number_list(motion, "axis", f"{path}.motion.axis", 3)
            if "range" in motion:
                _require_number_list(motion, "range", f"{path}.motion.range", 2)
            if "pivot" in motion:
                _require_number_list(motion, "pivot", f"{path}.motion.pivot", 3)
            for key in ["speed", "amplitude", "phase", "pitch"]:
                if key in motion and not _is_number(motion[key]):
                    raise ValueError(f"Invalid analysis payload at {path}.motion.{key}: expected a number")
        elif has_geometry:
            raise ValueError(f"Invalid analysis payload at {path}.motion: expected an object")

    return payload


def _require_string(payload: dict[str, Any], key: str, path: str) -> None:
    if not isinstance(payload.get(key), str) or not payload[key].strip():
        raise ValueError(f"Invalid analysis payload at {path}: expected a non-empty string")


def _require_string_list(payload: dict[str, Any], key: str, path: str) -> None:
    value = payload.get(key)
    if not isinstance(value, list) or not value or not all(isinstance(item, str) and item.strip() for item in value):
        raise ValueError(f"Invalid analysis payload at {path}: expected a non-empty list of strings")


def _require_number_list(payload: dict[str, Any], key: str, path: str, length: int) -> None:
    value = payload.get(key)
    if not isinstance(value, list) or len(value) != length or not all(_is_number(item) for item in value):
        raise ValueError(f"Invalid analysis payload at {path}: expected {length} numbers")


def _require_unit_number_list(payload: dict[str, Any], key: str, path: str, length: int) -> None:
    value = payload.get(key)
    if (
        not isinstance(value, list)
        or len(value) != length
        or not all(_is_number(item) and 0 <= item <= 1 for item in value)
    ):
        raise ValueError(f"Invalid analysis payload at {path}: expected {length} numbers from 0 to 1")


def _is_number(value: Any) -> bool:
    return isinstance(value, (int, float)) and not isinstance(value, bool)