"""Synthetic assembly graph generator for training data production.

Generates assembly graphs with known constraint classifications using
the pebble game and Jacobian verification. Each assembly is fully labeled
with per-constraint independence flags and assembly-level classification.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal

import numpy as np
from scipy.spatial.transform import Rotation

from solver.datagen.analysis import analyze_assembly
from solver.datagen.labeling import label_assembly
from solver.datagen.types import (
    ConstraintAnalysis,
    Joint,
    JointType,
    RigidBody,
)

if TYPE_CHECKING:
    from typing import Any

__all__ = [
    "COMPLEXITY_RANGES",
    "AxisStrategy",
    "ComplexityTier",
    "SyntheticAssemblyGenerator",
]

# ---------------------------------------------------------------------------
# Complexity tiers — ranges use exclusive upper bound for rng.integers()
# ---------------------------------------------------------------------------

ComplexityTier = Literal["simple", "medium", "complex"]

COMPLEXITY_RANGES: dict[str, tuple[int, int]] = {
    "simple": (2, 6),
    "medium": (6, 16),
    "complex": (16, 51),
}

# ---------------------------------------------------------------------------
# Axis sampling strategies
# ---------------------------------------------------------------------------

AxisStrategy = Literal["cardinal", "random", "near_parallel"]


class SyntheticAssemblyGenerator:
    """Generates assembly graphs with known minimal constraint sets.

    Uses the pebble game to incrementally build assemblies, tracking
    exactly which constraints are independent at each step. This produces
    labeled training data: (assembly_graph, constraint_set, labels).

    Labels per constraint:
        - independent: bool (does this constraint remove a DOF?)
        - redundant: bool (is this constraint overconstrained?)
        - minimal_set: bool (part of a minimal rigidity basis?)
    """

    def __init__(self, seed: int = 42) -> None:
        self.rng = np.random.default_rng(seed)

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _random_position(self, scale: float = 5.0) -> np.ndarray:
        """Generate random 3D position within [-scale, scale] cube."""
        return self.rng.uniform(-scale, scale, size=3)

    def _random_axis(self) -> np.ndarray:
        """Generate random normalized 3D axis."""
        axis = self.rng.standard_normal(3)
        axis /= np.linalg.norm(axis)
        return axis

    def _random_orientation(self) -> np.ndarray:
        """Generate a random 3x3 rotation matrix."""
        mat: np.ndarray = Rotation.random(random_state=self.rng).as_matrix()
        return mat

    def _cardinal_axis(self) -> np.ndarray:
        """Pick uniformly from the six signed cardinal directions."""
        axes = np.array(
            [
                [1, 0, 0],
                [-1, 0, 0],
                [0, 1, 0],
                [0, -1, 0],
                [0, 0, 1],
                [0, 0, -1],
            ],
            dtype=float,
        )
        result: np.ndarray = axes[int(self.rng.integers(6))]
        return result

    def _near_parallel_axis(
        self,
        base_axis: np.ndarray,
        perturbation_scale: float = 0.05,
    ) -> np.ndarray:
        """Return *base_axis* with a small random perturbation, re-normalized."""
        perturbed = base_axis + self.rng.standard_normal(3) * perturbation_scale
        return perturbed / np.linalg.norm(perturbed)

    def _sample_axis(self, strategy: AxisStrategy = "random") -> np.ndarray:
        """Sample a joint axis using the specified strategy."""
        if strategy == "cardinal":
            return self._cardinal_axis()
        if strategy == "near_parallel":
            return self._near_parallel_axis(np.array([0.0, 0.0, 1.0]))
        return self._random_axis()

    def _resolve_axis(
        self,
        strategy: AxisStrategy,
        parallel_axis_prob: float,
        shared_axis: np.ndarray | None,
    ) -> tuple[np.ndarray, np.ndarray | None]:
        """Return (axis_for_this_joint, shared_axis_to_propagate).

        On the first call where *shared_axis* is ``None`` and parallel
        injection triggers, a base axis is chosen and returned as
        *shared_axis* for subsequent calls.
        """
        if shared_axis is not None:
            return self._near_parallel_axis(shared_axis), shared_axis
        if parallel_axis_prob > 0 and self.rng.random() < parallel_axis_prob:
            base = self._sample_axis(strategy)
            return base.copy(), base
        return self._sample_axis(strategy), None

    def _select_joint_type(
        self,
        joint_types: JointType | list[JointType],
    ) -> JointType:
        """Select a joint type from a single type or list."""
        if isinstance(joint_types, list):
            idx = int(self.rng.integers(0, len(joint_types)))
            return joint_types[idx]
        return joint_types

    def _create_joint(
        self,
        joint_id: int,
        body_a_id: int,
        body_b_id: int,
        pos_a: np.ndarray,
        pos_b: np.ndarray,
        joint_type: JointType,
        *,
        axis: np.ndarray | None = None,
        orient_a: np.ndarray | None = None,
        orient_b: np.ndarray | None = None,
    ) -> Joint:
        """Create a joint between two bodies.

        When orientations are provided, anchor points are offset from
        each body's center along a random local direction rotated into
        world frame, rather than placed at the midpoint.
        """
        if orient_a is not None and orient_b is not None:
            dist = max(float(np.linalg.norm(pos_b - pos_a)), 0.1)
            offset_scale = dist * 0.2
            local_a = self.rng.standard_normal(3) * offset_scale
            local_b = self.rng.standard_normal(3) * offset_scale
            anchor_a = pos_a + orient_a @ local_a
            anchor_b = pos_b + orient_b @ local_b
        else:
            anchor = (pos_a + pos_b) / 2.0
            anchor_a = anchor
            anchor_b = anchor

        return Joint(
            joint_id=joint_id,
            body_a=body_a_id,
            body_b=body_b_id,
            joint_type=joint_type,
            anchor_a=anchor_a,
            anchor_b=anchor_b,
            axis=axis if axis is not None else self._random_axis(),
        )

    # ------------------------------------------------------------------
    # Original generators (chain / rigid / overconstrained)
    # ------------------------------------------------------------------

    def generate_chain_assembly(
        self,
        n_bodies: int,
        joint_type: JointType = JointType.REVOLUTE,
        *,
        grounded: bool = True,
        axis_strategy: AxisStrategy = "random",
        parallel_axis_prob: float = 0.0,
    ) -> tuple[list[RigidBody], list[Joint], ConstraintAnalysis]:
        """Generate a serial kinematic chain.

        Simple but useful: each body connects to the next with the
        specified joint type. Results in an underconstrained assembly
        (serial chain is never rigid without closing loops).
        """
        bodies = []
        joints = []

        for i in range(n_bodies):
            pos = np.array([i * 2.0, 0.0, 0.0])
            bodies.append(
                RigidBody(
                    body_id=i,
                    position=pos,
                    orientation=self._random_orientation(),
                )
            )

        shared_axis: np.ndarray | None = None
        for i in range(n_bodies - 1):
            axis, shared_axis = self._resolve_axis(
                axis_strategy,
                parallel_axis_prob,
                shared_axis,
            )
            joints.append(
                self._create_joint(
                    i,
                    i,
                    i + 1,
                    bodies[i].position,
                    bodies[i + 1].position,
                    joint_type,
                    axis=axis,
                    orient_a=bodies[i].orientation,
                    orient_b=bodies[i + 1].orientation,
                )
            )

        analysis = analyze_assembly(
            bodies,
            joints,
            ground_body=0 if grounded else None,
        )
        return bodies, joints, analysis

    def generate_rigid_assembly(
        self,
        n_bodies: int,
        *,
        grounded: bool = True,
        axis_strategy: AxisStrategy = "random",
        parallel_axis_prob: float = 0.0,
    ) -> tuple[list[RigidBody], list[Joint], ConstraintAnalysis]:
        """Generate a minimally rigid assembly by adding joints until rigid.

        Strategy: start with fixed joints on a spanning tree (guarantees
        rigidity), then randomly relax some to weaker joint types while
        maintaining rigidity via the pebble game check.
        """
        bodies = []
        for i in range(n_bodies):
            bodies.append(
                RigidBody(
                    body_id=i,
                    position=self._random_position(),
                    orientation=self._random_orientation(),
                )
            )

        # Build spanning tree with fixed joints (overconstrained)
        joints: list[Joint] = []
        shared_axis: np.ndarray | None = None
        for i in range(1, n_bodies):
            parent = int(self.rng.integers(0, i))
            axis, shared_axis = self._resolve_axis(
                axis_strategy,
                parallel_axis_prob,
                shared_axis,
            )
            joints.append(
                self._create_joint(
                    i - 1,
                    parent,
                    i,
                    bodies[parent].position,
                    bodies[i].position,
                    JointType.FIXED,
                    axis=axis,
                    orient_a=bodies[parent].orientation,
                    orient_b=bodies[i].orientation,
                )
            )

        # Try relaxing joints to weaker types while maintaining rigidity
        weaker_types = [
            JointType.REVOLUTE,
            JointType.CYLINDRICAL,
            JointType.BALL,
        ]

        ground = 0 if grounded else None
        for idx in self.rng.permutation(len(joints)):
            original_type = joints[idx].joint_type
            for candidate in weaker_types:
                joints[idx].joint_type = candidate
                analysis = analyze_assembly(bodies, joints, ground_body=ground)
                if analysis.is_rigid:
                    break  # Keep the weaker type
            else:
                joints[idx].joint_type = original_type

        analysis = analyze_assembly(bodies, joints, ground_body=ground)
        return bodies, joints, analysis

    def generate_overconstrained_assembly(
        self,
        n_bodies: int,
        extra_joints: int = 2,
        *,
        grounded: bool = True,
        axis_strategy: AxisStrategy = "random",
        parallel_axis_prob: float = 0.0,
    ) -> tuple[list[RigidBody], list[Joint], ConstraintAnalysis]:
        """Generate an assembly with known redundant constraints.

        Starts with a rigid assembly, then adds extra joints that
        the pebble game will flag as redundant.
        """
        bodies, joints, _ = self.generate_rigid_assembly(
            n_bodies,
            grounded=grounded,
            axis_strategy=axis_strategy,
            parallel_axis_prob=parallel_axis_prob,
        )

        joint_id = len(joints)
        shared_axis: np.ndarray | None = None
        for _ in range(extra_joints):
            a, b = self.rng.choice(n_bodies, size=2, replace=False)
            _overcon_types = [
                JointType.REVOLUTE,
                JointType.FIXED,
                JointType.BALL,
            ]
            jtype = _overcon_types[int(self.rng.integers(len(_overcon_types)))]
            axis, shared_axis = self._resolve_axis(
                axis_strategy,
                parallel_axis_prob,
                shared_axis,
            )
            joints.append(
                self._create_joint(
                    joint_id,
                    int(a),
                    int(b),
                    bodies[int(a)].position,
                    bodies[int(b)].position,
                    jtype,
                    axis=axis,
                    orient_a=bodies[int(a)].orientation,
                    orient_b=bodies[int(b)].orientation,
                )
            )
            joint_id += 1

        ground = 0 if grounded else None
        analysis = analyze_assembly(bodies, joints, ground_body=ground)
        return bodies, joints, analysis

    # ------------------------------------------------------------------
    # New topology generators
    # ------------------------------------------------------------------

    def generate_tree_assembly(
        self,
        n_bodies: int,
        joint_types: JointType | list[JointType] = JointType.REVOLUTE,
        branching_factor: int = 3,
        *,
        grounded: bool = True,
        axis_strategy: AxisStrategy = "random",
        parallel_axis_prob: float = 0.0,
    ) -> tuple[list[RigidBody], list[Joint], ConstraintAnalysis]:
        """Generate a random tree topology with configurable branching.

        Creates a tree where each body can have up to *branching_factor*
        children. Different branches can use different joint types if a
        list is provided. Always underconstrained (no closed loops).

        Args:
            n_bodies: Total bodies (root + children).
            joint_types: Single type or list to sample from per joint.
            branching_factor: Max children per parent (1-5 recommended).
        """
        bodies: list[RigidBody] = [
            RigidBody(
                body_id=0,
                position=np.zeros(3),
                orientation=self._random_orientation(),
            )
        ]
        joints: list[Joint] = []

        available_parents = [0]
        next_id = 1
        joint_id = 0
        shared_axis: np.ndarray | None = None

        while next_id < n_bodies and available_parents:
            pidx = int(self.rng.integers(0, len(available_parents)))
            parent_id = available_parents[pidx]
            parent_pos = bodies[parent_id].position

            max_children = min(branching_factor, n_bodies - next_id)
            n_children = int(self.rng.integers(1, max_children + 1))

            for _ in range(n_children):
                direction = self._random_axis()
                distance = self.rng.uniform(1.5, 3.0)
                child_pos = parent_pos + direction * distance
                child_orient = self._random_orientation()

                bodies.append(
                    RigidBody(
                        body_id=next_id,
                        position=child_pos,
                        orientation=child_orient,
                    )
                )
                jtype = self._select_joint_type(joint_types)
                axis, shared_axis = self._resolve_axis(
                    axis_strategy,
                    parallel_axis_prob,
                    shared_axis,
                )
                joints.append(
                    self._create_joint(
                        joint_id,
                        parent_id,
                        next_id,
                        parent_pos,
                        child_pos,
                        jtype,
                        axis=axis,
                        orient_a=bodies[parent_id].orientation,
                        orient_b=child_orient,
                    )
                )

                available_parents.append(next_id)
                next_id += 1
                joint_id += 1
                if next_id >= n_bodies:
                    break

            # Retire parent if it reached branching limit or randomly
            if n_children >= branching_factor or self.rng.random() < 0.3:
                available_parents.pop(pidx)

        analysis = analyze_assembly(
            bodies,
            joints,
            ground_body=0 if grounded else None,
        )
        return bodies, joints, analysis

    def generate_loop_assembly(
        self,
        n_bodies: int,
        joint_types: JointType | list[JointType] = JointType.REVOLUTE,
        *,
        grounded: bool = True,
        axis_strategy: AxisStrategy = "random",
        parallel_axis_prob: float = 0.0,
    ) -> tuple[list[RigidBody], list[Joint], ConstraintAnalysis]:
        """Generate a single closed loop (ring) of bodies.

        The closing constraint introduces redundancy, making this
        useful for generating overconstrained training data.

        Args:
            n_bodies: Bodies in the ring (>= 3).
            joint_types: Single type or list to sample from per joint.

        Raises:
            ValueError: If *n_bodies* < 3.
        """
        if n_bodies < 3:
            msg = "Loop assembly requires at least 3 bodies"
            raise ValueError(msg)

        bodies: list[RigidBody] = []
        joints: list[Joint] = []

        base_radius = max(2.0, n_bodies * 0.4)
        for i in range(n_bodies):
            angle = 2 * np.pi * i / n_bodies
            radius = base_radius + self.rng.uniform(-0.5, 0.5)
            x = radius * np.cos(angle)
            y = radius * np.sin(angle)
            z = float(self.rng.uniform(-0.2, 0.2))
            bodies.append(
                RigidBody(
                    body_id=i,
                    position=np.array([x, y, z]),
                    orientation=self._random_orientation(),
                )
            )

        shared_axis: np.ndarray | None = None
        for i in range(n_bodies):
            next_i = (i + 1) % n_bodies
            jtype = self._select_joint_type(joint_types)
            axis, shared_axis = self._resolve_axis(
                axis_strategy,
                parallel_axis_prob,
                shared_axis,
            )
            joints.append(
                self._create_joint(
                    i,
                    i,
                    next_i,
                    bodies[i].position,
                    bodies[next_i].position,
                    jtype,
                    axis=axis,
                    orient_a=bodies[i].orientation,
                    orient_b=bodies[next_i].orientation,
                )
            )

        analysis = analyze_assembly(
            bodies,
            joints,
            ground_body=0 if grounded else None,
        )
        return bodies, joints, analysis

    def generate_star_assembly(
        self,
        n_bodies: int,
        joint_types: JointType | list[JointType] = JointType.REVOLUTE,
        *,
        grounded: bool = True,
        axis_strategy: AxisStrategy = "random",
        parallel_axis_prob: float = 0.0,
    ) -> tuple[list[RigidBody], list[Joint], ConstraintAnalysis]:
        """Generate a star topology with central hub and satellites.

        Body 0 is the hub; all other bodies connect directly to it.
        Underconstrained because there are no inter-satellite connections.

        Args:
            n_bodies: Total bodies including hub (>= 2).
            joint_types: Single type or list to sample from per joint.

        Raises:
            ValueError: If *n_bodies* < 2.
        """
        if n_bodies < 2:
            msg = "Star assembly requires at least 2 bodies"
            raise ValueError(msg)

        hub_orient = self._random_orientation()
        bodies: list[RigidBody] = [
            RigidBody(
                body_id=0,
                position=np.zeros(3),
                orientation=hub_orient,
            )
        ]
        joints: list[Joint] = []

        shared_axis: np.ndarray | None = None
        for i in range(1, n_bodies):
            direction = self._random_axis()
            distance = self.rng.uniform(2.0, 5.0)
            pos = direction * distance
            sat_orient = self._random_orientation()
            bodies.append(RigidBody(body_id=i, position=pos, orientation=sat_orient))

            jtype = self._select_joint_type(joint_types)
            axis, shared_axis = self._resolve_axis(
                axis_strategy,
                parallel_axis_prob,
                shared_axis,
            )
            joints.append(
                self._create_joint(
                    i - 1,
                    0,
                    i,
                    np.zeros(3),
                    pos,
                    jtype,
                    axis=axis,
                    orient_a=hub_orient,
                    orient_b=sat_orient,
                )
            )

        analysis = analyze_assembly(
            bodies,
            joints,
            ground_body=0 if grounded else None,
        )
        return bodies, joints, analysis

    def generate_mixed_assembly(
        self,
        n_bodies: int,
        joint_types: JointType | list[JointType] = JointType.REVOLUTE,
        edge_density: float = 0.3,
        *,
        grounded: bool = True,
        axis_strategy: AxisStrategy = "random",
        parallel_axis_prob: float = 0.0,
    ) -> tuple[list[RigidBody], list[Joint], ConstraintAnalysis]:
        """Generate a mixed topology combining tree and loop elements.

        Builds a spanning tree for connectivity, then adds extra edges
        based on *edge_density* to create loops and redundancy.

        Args:
            n_bodies: Number of bodies.
            joint_types: Single type or list to sample from per joint.
            edge_density: Fraction of non-tree edges to add (0.0-1.0).

        Raises:
            ValueError: If *edge_density* not in [0.0, 1.0].
        """
        if not 0.0 <= edge_density <= 1.0:
            msg = "edge_density must be in [0.0, 1.0]"
            raise ValueError(msg)

        bodies: list[RigidBody] = []
        joints: list[Joint] = []

        for i in range(n_bodies):
            bodies.append(
                RigidBody(
                    body_id=i,
                    position=self._random_position(),
                    orientation=self._random_orientation(),
                )
            )

        # Phase 1: spanning tree
        joint_id = 0
        existing_edges: set[frozenset[int]] = set()
        shared_axis: np.ndarray | None = None
        for i in range(1, n_bodies):
            parent = int(self.rng.integers(0, i))
            jtype = self._select_joint_type(joint_types)
            axis, shared_axis = self._resolve_axis(
                axis_strategy,
                parallel_axis_prob,
                shared_axis,
            )
            joints.append(
                self._create_joint(
                    joint_id,
                    parent,
                    i,
                    bodies[parent].position,
                    bodies[i].position,
                    jtype,
                    axis=axis,
                    orient_a=bodies[parent].orientation,
                    orient_b=bodies[i].orientation,
                )
            )
            existing_edges.add(frozenset([parent, i]))
            joint_id += 1

        # Phase 2: extra edges based on density
        candidates: list[tuple[int, int]] = []
        for i in range(n_bodies):
            for j in range(i + 1, n_bodies):
                if frozenset([i, j]) not in existing_edges:
                    candidates.append((i, j))

        n_extra = int(edge_density * len(candidates))
        self.rng.shuffle(candidates)

        for a, b in candidates[:n_extra]:
            jtype = self._select_joint_type(joint_types)
            axis, shared_axis = self._resolve_axis(
                axis_strategy,
                parallel_axis_prob,
                shared_axis,
            )
            joints.append(
                self._create_joint(
                    joint_id,
                    a,
                    b,
                    bodies[a].position,
                    bodies[b].position,
                    jtype,
                    axis=axis,
                    orient_a=bodies[a].orientation,
                    orient_b=bodies[b].orientation,
                )
            )
            joint_id += 1

        analysis = analyze_assembly(
            bodies,
            joints,
            ground_body=0 if grounded else None,
        )
        return bodies, joints, analysis

    # ------------------------------------------------------------------
    # Batch generation
    # ------------------------------------------------------------------

    def generate_training_batch(
        self,
        batch_size: int = 100,
        n_bodies_range: tuple[int, int] | None = None,
        complexity_tier: ComplexityTier | None = None,
        *,
        axis_strategy: AxisStrategy = "random",
        parallel_axis_prob: float = 0.0,
        grounded_ratio: float = 1.0,
    ) -> list[dict[str, Any]]:
        """Generate a batch of labeled training examples.

        Each example contains body positions, joint descriptions,
        per-joint independence labels, and assembly-level classification.

        Args:
            batch_size: Number of assemblies to generate.
            n_bodies_range: ``(min, max_exclusive)`` body count.
                Overridden by *complexity_tier* when both are given.
            complexity_tier: Predefined range (``"simple"`` / ``"medium"``
                / ``"complex"``).  Overrides *n_bodies_range*.
            axis_strategy: Axis sampling strategy for joint axes.
            parallel_axis_prob: Probability of parallel axis injection.
            grounded_ratio: Fraction of examples that are grounded.
        """
        if complexity_tier is not None:
            n_bodies_range = COMPLEXITY_RANGES[complexity_tier]
        elif n_bodies_range is None:
            n_bodies_range = (3, 8)

        _joint_pool = [
            JointType.REVOLUTE,
            JointType.BALL,
            JointType.CYLINDRICAL,
            JointType.FIXED,
        ]

        geo_kw: dict[str, Any] = {
            "axis_strategy": axis_strategy,
            "parallel_axis_prob": parallel_axis_prob,
        }

        examples: list[dict[str, Any]] = []

        for i in range(batch_size):
            n = int(self.rng.integers(*n_bodies_range))
            gen_idx = int(self.rng.integers(7))
            grounded = bool(self.rng.random() < grounded_ratio)

            if gen_idx == 0:
                _chain_types = [
                    JointType.REVOLUTE,
                    JointType.BALL,
                    JointType.CYLINDRICAL,
                ]
                jtype = _chain_types[int(self.rng.integers(len(_chain_types)))]
                bodies, joints, analysis = self.generate_chain_assembly(
                    n,
                    jtype,
                    grounded=grounded,
                    **geo_kw,
                )
                gen_name = "chain"
            elif gen_idx == 1:
                bodies, joints, analysis = self.generate_rigid_assembly(
                    n,
                    grounded=grounded,
                    **geo_kw,
                )
                gen_name = "rigid"
            elif gen_idx == 2:
                extra = int(self.rng.integers(1, 4))
                bodies, joints, analysis = self.generate_overconstrained_assembly(
                    n,
                    extra,
                    grounded=grounded,
                    **geo_kw,
                )
                gen_name = "overconstrained"
            elif gen_idx == 3:
                branching = int(self.rng.integers(2, 5))
                bodies, joints, analysis = self.generate_tree_assembly(
                    n,
                    _joint_pool,
                    branching,
                    grounded=grounded,
                    **geo_kw,
                )
                gen_name = "tree"
            elif gen_idx == 4:
                n = max(n, 3)
                bodies, joints, analysis = self.generate_loop_assembly(
                    n,
                    _joint_pool,
                    grounded=grounded,
                    **geo_kw,
                )
                gen_name = "loop"
            elif gen_idx == 5:
                n = max(n, 2)
                bodies, joints, analysis = self.generate_star_assembly(
                    n,
                    _joint_pool,
                    grounded=grounded,
                    **geo_kw,
                )
                gen_name = "star"
            else:
                density = float(self.rng.uniform(0.2, 0.5))
                bodies, joints, analysis = self.generate_mixed_assembly(
                    n,
                    _joint_pool,
                    density,
                    grounded=grounded,
                    **geo_kw,
                )
                gen_name = "mixed"

            # Produce ground truth labels (includes ConstraintAnalysis)
            ground = 0 if grounded else None
            labels = label_assembly(bodies, joints, ground_body=ground)
            analysis = labels.analysis

            # Build per-joint labels from edge results
            joint_labels: dict[int, dict[str, int]] = {}
            for result in analysis.per_edge_results:
                jid = result["joint_id"]
                if jid not in joint_labels:
                    joint_labels[jid] = {
                        "independent_constraints": 0,
                        "redundant_constraints": 0,
                        "total_constraints": 0,
                    }
                joint_labels[jid]["total_constraints"] += 1
                if result["independent"]:
                    joint_labels[jid]["independent_constraints"] += 1
                else:
                    joint_labels[jid]["redundant_constraints"] += 1

            examples.append(
                {
                    "example_id": i,
                    "generator_type": gen_name,
                    "grounded": grounded,
                    "n_bodies": len(bodies),
                    "n_joints": len(joints),
                    "body_positions": [b.position.tolist() for b in bodies],
                    "body_orientations": [b.orientation.tolist() for b in bodies],
                    "joints": [
                        {
                            "joint_id": j.joint_id,
                            "body_a": j.body_a,
                            "body_b": j.body_b,
                            "type": j.joint_type.name,
                            "axis": j.axis.tolist(),
                        }
                        for j in joints
                    ],
                    "joint_labels": joint_labels,
                    "labels": labels.to_dict(),
                    "assembly_classification": (analysis.combinatorial_classification),
                    "is_rigid": analysis.is_rigid,
                    "is_minimally_rigid": analysis.is_minimally_rigid,
                    "internal_dof": analysis.jacobian_internal_dof,
                    "geometric_degeneracies": (analysis.geometric_degeneracies),
                }
            )

        return examples
