"""Tests for graph decomposition and cluster-by-cluster solving."""

import math

import pytest
from kindred_solver.constraints import (
    CoincidentConstraint,
    CylindricalConstraint,
    FixedConstraint,
    ParallelConstraint,
    RevoluteConstraint,
    SliderConstraint,
)
from kindred_solver.decompose import (
    SolveCluster,
    _constraints_for_bodies,
    build_constraint_graph_simple,
    build_solve_order,
    classify_cluster_rigidity,
    find_clusters,
    residual_count_by_name,
    solve_decomposed,
)
from kindred_solver.dof import count_dof
from kindred_solver.entities import RigidBody
from kindred_solver.newton import newton_solve
from kindred_solver.params import ParamTable
from kindred_solver.prepass import single_equation_pass, substitution_pass

ID_QUAT = (1, 0, 0, 0)
c45 = math.cos(math.pi / 4)
s45 = math.sin(math.pi / 4)
ROT_90Z = (c45, 0, 0, s45)


# ============================================================================
# DOF table tests
# ============================================================================


class TestResidualCount:
    def test_known_types(self):
        assert residual_count_by_name("Fixed") == 6
        assert residual_count_by_name("Revolute") == 5
        assert residual_count_by_name("Cylindrical") == 4
        assert residual_count_by_name("Ball") == 3
        assert residual_count_by_name("Coincident") == 3
        assert residual_count_by_name("Parallel") == 2
        assert residual_count_by_name("Perpendicular") == 1
        assert residual_count_by_name("DistancePointPoint") == 1

    def test_stubs_zero(self):
        assert residual_count_by_name("Cam") == 0
        assert residual_count_by_name("Slot") == 0

    def test_unknown_zero(self):
        assert residual_count_by_name("DoesNotExist") == 0


# ============================================================================
# Graph construction tests
# ============================================================================


class TestBuildConstraintGraph:
    def test_simple_pair(self):
        """Two bodies, one Revolute constraint."""
        G = build_constraint_graph_simple(
            [("A", "B", "Revolute", 0)],
            grounded={"A"},
        )
        assert set(G.nodes()) == {"A", "B"}
        assert G.number_of_edges() == 1
        assert G.nodes["A"]["grounded"] is True
        assert G.nodes["B"]["grounded"] is False

    def test_chain(self):
        """Chain: A-B-C-D."""
        G = build_constraint_graph_simple(
            [
                ("A", "B", "Revolute", 0),
                ("B", "C", "Revolute", 1),
                ("C", "D", "Revolute", 2),
            ],
        )
        assert G.number_of_nodes() == 4
        assert G.number_of_edges() == 3

    def test_multi_edge(self):
        """Two constraints between same body pair → two edges."""
        G = build_constraint_graph_simple(
            [
                ("A", "B", "Revolute", 0),
                ("A", "B", "Parallel", 1),
            ],
        )
        assert G.number_of_nodes() == 2
        assert G.number_of_edges() == 2

    def test_stub_excluded(self):
        """Stub constraints (Cam, Slot) excluded from graph."""
        G = build_constraint_graph_simple(
            [
                ("A", "B", "Revolute", 0),
                ("B", "C", "Cam", 1),
            ],
        )
        assert G.number_of_nodes() == 2  # C not added
        assert G.number_of_edges() == 1


# ============================================================================
# Biconnected decomposition tests
# ============================================================================


class TestFindClusters:
    def test_chain_decomposes(self):
        """Chain A-B-C-D: 3 biconnected components, 2 articulation points."""
        G = build_constraint_graph_simple(
            [
                ("A", "B", "Revolute", 0),
                ("B", "C", "Revolute", 1),
                ("C", "D", "Revolute", 2),
            ],
        )
        clusters, artic = find_clusters(G)
        assert len(clusters) == 3
        assert artic == {"B", "C"}

    def test_loop_single_cluster(self):
        """Loop A-B-C-A: single biconnected component, no articulation points."""
        G = build_constraint_graph_simple(
            [
                ("A", "B", "Revolute", 0),
                ("B", "C", "Revolute", 1),
                ("C", "A", "Revolute", 2),
            ],
        )
        clusters, artic = find_clusters(G)
        assert len(clusters) == 1
        assert artic == set()

    def test_star_decomposes(self):
        """Star: center connected to 3 leaves.

        3 biconnected components (each is center + one leaf).
        Center is the only articulation point.
        """
        G = build_constraint_graph_simple(
            [
                ("center", "L1", "Revolute", 0),
                ("center", "L2", "Revolute", 1),
                ("center", "L3", "Revolute", 2),
            ],
        )
        clusters, artic = find_clusters(G)
        assert len(clusters) == 3
        assert artic == {"center"}

    def test_single_edge(self):
        """Two nodes, one edge: single biconnected component."""
        G = build_constraint_graph_simple(
            [("A", "B", "Fixed", 0)],
        )
        clusters, artic = find_clusters(G)
        assert len(clusters) == 1
        assert artic == set()

    def test_tree_with_loop(self):
        """A-B-C-D with a loop B-C-E-B.

        The loop {B, C, E} forms one biconnected component.
        A-B and C-D are separate biconnected components.
        Articulation points: B and C.
        """
        G = build_constraint_graph_simple(
            [
                ("A", "B", "Revolute", 0),
                ("B", "C", "Revolute", 1),
                ("C", "E", "Revolute", 2),
                ("E", "B", "Revolute", 3),
                ("C", "D", "Revolute", 4),
            ],
        )
        clusters, artic = find_clusters(G)
        # 3 clusters: {A,B}, {B,C,E}, {C,D}
        assert len(clusters) == 3
        assert artic == {"B", "C"}


# ============================================================================
# Solve order tests
# ============================================================================


class TestBuildSolveOrder:
    def test_chain_grounded_at_start(self):
        """Chain: ground-A-B-C.

        Root cluster is {ground,A}, solved first. Then {A,B}, then {B,C}.
        """
        G = build_constraint_graph_simple(
            [
                ("ground", "A", "Revolute", 0),
                ("A", "B", "Revolute", 1),
                ("B", "C", "Revolute", 2),
            ],
            grounded={"ground"},
        )
        clusters, artic = find_clusters(G)
        solve_order = build_solve_order(G, clusters, artic, {"ground"})

        assert len(solve_order) == 3
        # First cluster should contain ground (root)
        assert "ground" in solve_order[0].bodies
        assert solve_order[0].has_ground is True
        # Last cluster should be the leaf
        assert "C" in solve_order[-1].bodies

    def test_star_grounded_center(self):
        """Star with grounded center: root cluster first, then leaves."""
        G = build_constraint_graph_simple(
            [
                ("center", "L1", "Revolute", 0),
                ("center", "L2", "Revolute", 1),
                ("center", "L3", "Revolute", 2),
            ],
            grounded={"center"},
        )
        clusters, artic = find_clusters(G)
        solve_order = build_solve_order(G, clusters, artic, {"center"})

        assert len(solve_order) == 3
        # First cluster contains grounded center
        assert "center" in solve_order[0].bodies
        assert solve_order[0].has_ground is True
        # All clusters have center as boundary
        for sc in solve_order:
            assert "center" in sc.boundary_bodies or "center" in sc.bodies

    def test_single_cluster_no_boundary(self):
        """Single cluster: no boundary bodies."""
        G = build_constraint_graph_simple(
            [("A", "B", "Revolute", 0)],
            grounded={"A"},
        )
        clusters, artic = find_clusters(G)
        solve_order = build_solve_order(G, clusters, artic, {"A"})

        assert len(solve_order) == 1
        assert solve_order[0].boundary_bodies == set()
        assert solve_order[0].has_ground is True

    def test_constraint_assignment(self):
        """Constraints are correctly assigned to clusters."""
        G = build_constraint_graph_simple(
            [
                ("A", "B", "Revolute", 0),
                ("B", "C", "Revolute", 1),
            ],
        )
        clusters, artic = find_clusters(G)
        solve_order = build_solve_order(G, clusters, artic, set())

        # Each cluster should have exactly 1 constraint
        all_indices = []
        for sc in solve_order:
            all_indices.extend(sc.constraint_indices)
        assert sorted(all_indices) == [0, 1]


# ============================================================================
# Integration: cluster solving with actual Newton solver
# ============================================================================


def _monolithic_solve(pt, all_bodies, constraint_objs):
    """Monolithic solve for comparison."""
    residuals = []
    for c in constraint_objs:
        residuals.extend(c.residuals())
    quat_groups = []
    for body in all_bodies:
        if not body.grounded:
            residuals.append(body.quat_norm_residual())
            quat_groups.append(body.quat_param_names())
    residuals = substitution_pass(residuals, pt)
    residuals = single_equation_pass(residuals, pt)
    return newton_solve(residuals, pt, quat_groups=quat_groups, max_iter=100, tol=1e-10)


class TestClusterSolve:
    def test_chain_solve_matches_monolithic(self):
        """Chain: ground → A → B.  Decomposed solve matches monolithic.

        Two biconnected components: {ground,A} and {A,B}.
        A is the articulation point.
        """
        # --- Decomposed solve ---
        pt = ParamTable()
        ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True)
        body_a = RigidBody("A", pt, (3, 2, 0), ID_QUAT)
        body_b = RigidBody("B", pt, (8, 3, 0), ID_QUAT)
        bodies = {"ground": ground, "A": body_a, "B": body_b}

        c0 = RevoluteConstraint(ground, (0, 0, 0), ID_QUAT, body_a, (0, 0, 0), ID_QUAT)
        c1 = RevoluteConstraint(body_a, (5, 0, 0), ID_QUAT, body_b, (0, 0, 0), ID_QUAT)
        constraint_objs = [c0, c1]
        constraint_indices = [0, 1]

        # Solve order: grounded cluster first, then outward.
        # {ground,A} first (A gets solved), then {A,B} with A as boundary.
        clusters = [
            SolveCluster(
                bodies={"ground", "A"},
                constraint_indices=[0],
                boundary_bodies={"A"},
                has_ground=True,
            ),
            SolveCluster(
                bodies={"A", "B"},
                constraint_indices=[1],
                boundary_bodies={"A"},
                has_ground=False,
            ),
        ]

        converged = solve_decomposed(
            clusters,
            bodies,
            constraint_objs,
            constraint_indices,
            pt,
        )
        assert converged

        env = pt.get_env()
        pos_a = body_a.extract_position(env)
        pos_b = body_b.extract_position(env)

        # A should be at origin (revolute to ground origin)
        assert abs(pos_a[0]) < 1e-8
        assert abs(pos_a[1]) < 1e-8
        assert abs(pos_a[2]) < 1e-8

        # B should be at (5,0,0) — revolute at A's (5,0,0) marker, B's origin
        assert abs(pos_b[0] - 5.0) < 1e-8
        assert abs(pos_b[1]) < 1e-8
        assert abs(pos_b[2]) < 1e-8

    def test_star_solve(self):
        """Star: grounded center with 3 revolute arms.

        Each arm is a separate cluster. Center is articulation point.
        Solve center-containing clusters first, then arms.
        """
        pt = ParamTable()
        center = RigidBody("C", pt, (0, 0, 0), ID_QUAT, grounded=True)
        arm1 = RigidBody("L1", pt, (5, 5, 0), ID_QUAT)
        arm2 = RigidBody("L2", pt, (5, -5, 0), ID_QUAT)
        arm3 = RigidBody("L3", pt, (-5, 5, 0), ID_QUAT)
        bodies = {"C": center, "L1": arm1, "L2": arm2, "L3": arm3}

        c0 = RevoluteConstraint(center, (3, 0, 0), ID_QUAT, arm1, (0, 0, 0), ID_QUAT)
        c1 = RevoluteConstraint(center, (0, 3, 0), ID_QUAT, arm2, (0, 0, 0), ID_QUAT)
        c2 = RevoluteConstraint(center, (-3, 0, 0), ID_QUAT, arm3, (0, 0, 0), ID_QUAT)
        constraint_objs = [c0, c1, c2]
        constraint_indices = [0, 1, 2]

        # Each arm cluster has center as boundary. Since center is grounded,
        # its params are already fixed. Solve order doesn't matter much.
        clusters = [
            SolveCluster(
                bodies={"C", "L1"},
                constraint_indices=[0],
                boundary_bodies={"C"},
                has_ground=True,
            ),
            SolveCluster(
                bodies={"C", "L2"},
                constraint_indices=[1],
                boundary_bodies={"C"},
                has_ground=True,
            ),
            SolveCluster(
                bodies={"C", "L3"},
                constraint_indices=[2],
                boundary_bodies={"C"},
                has_ground=True,
            ),
        ]

        converged = solve_decomposed(
            clusters,
            bodies,
            constraint_objs,
            constraint_indices,
            pt,
        )
        assert converged

        env = pt.get_env()
        # Each arm should be at its marker position
        assert abs(arm1.extract_position(env)[0] - 3.0) < 1e-8
        assert abs(arm2.extract_position(env)[1] - 3.0) < 1e-8
        assert abs(arm3.extract_position(env)[0] + 3.0) < 1e-8

    def test_single_cluster_solve(self):
        """Single cluster: decomposed solve behaves same as monolithic."""
        pt = ParamTable()
        ground = RigidBody("g", pt, (0, 0, 0), ID_QUAT, grounded=True)
        body = RigidBody("b", pt, (5, 5, 5), ID_QUAT)
        bodies = {"g": ground, "b": body}

        c0 = RevoluteConstraint(ground, (0, 0, 0), ID_QUAT, body, (0, 0, 0), ID_QUAT)
        constraint_objs = [c0]
        constraint_indices = [0]

        clusters = [
            SolveCluster(
                bodies={"g", "b"},
                constraint_indices=[0],
                boundary_bodies=set(),
                has_ground=True,
            ),
        ]

        converged = solve_decomposed(
            clusters,
            bodies,
            constraint_objs,
            constraint_indices,
            pt,
        )
        assert converged

        env = pt.get_env()
        pos = body.extract_position(env)
        assert abs(pos[0]) < 1e-8
        assert abs(pos[1]) < 1e-8
        assert abs(pos[2]) < 1e-8


# ============================================================================
# Full decompose() pipeline tests
# ============================================================================


class TestDecomposePipeline:
    """Test the full decompose() entry point (requires kcsolve for
    build_constraint_graph, so we test via build_constraint_graph_simple
    and manual cluster construction instead)."""

    def test_chain_topology(self):
        """Chain decomposition produces correct number of clusters."""
        G = build_constraint_graph_simple(
            [
                ("g", "A", "Revolute", 0),
                ("A", "B", "Revolute", 1),
                ("B", "C", "Revolute", 2),
                ("C", "D", "Revolute", 3),
            ],
            grounded={"g"},
        )
        clusters, artic = find_clusters(G)
        solve_order = build_solve_order(G, clusters, artic, {"g"})

        # Chain of 5 nodes → 4 biconnected components
        assert len(solve_order) == 4
        # First cluster should contain ground
        assert "g" in solve_order[0].bodies

    def test_disconnected_components(self):
        """Two disconnected sub-assemblies produce independent cluster groups."""
        G = build_constraint_graph_simple(
            [
                ("A", "B", "Revolute", 0),
                ("C", "D", "Revolute", 1),
            ],
        )
        # Two connected components, each single-cluster
        components = list(nx.connected_components(G))
        assert len(components) == 2

        # Each component decomposes to 1 cluster
        for comp_nodes in components:
            sub = G.subgraph(comp_nodes).copy()
            clusters, artic = find_clusters(sub)
            assert len(clusters) == 1
            assert artic == set()


import networkx as nx
from kindred_solver.bfgs import bfgs_solve

# ============================================================================
# Integration: decomposed solve vs monolithic on larger assemblies
# ============================================================================


class TestDecomposedVsMonolithic:
    """Verify decomposed solving produces equivalent results to monolithic."""

    def _build_chain(self, n_bodies, pt, displaced=True):
        """Build a revolute chain: ground → B0 → B1 → ... → B(n-1).

        Each body is connected at (2,0,0) local → (0,0,0) local of next.
        If displaced, bodies start at slightly wrong positions.
        Returns (bodies_dict, constraint_objs, constraint_indices).
        """
        ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True)
        bodies = {"ground": ground}
        constraints = []
        indices = []

        prev_body = ground
        for i in range(n_bodies):
            name = f"B{i}"
            # Correct position would be ((i+1)*2, 0, 0), displace slightly
            x = (i + 1) * 2 + (0.5 if displaced else 0)
            y = 0.3 if displaced else 0
            body = RigidBody(name, pt, (x, y, 0), ID_QUAT)
            bodies[name] = body

            c = RevoluteConstraint(
                prev_body,
                (2, 0, 0),
                ID_QUAT,
                body,
                (0, 0, 0),
                ID_QUAT,
            )
            constraints.append(c)
            indices.append(i)
            prev_body = body

        return bodies, constraints, indices

    def test_chain_10_bodies(self):
        """10-body revolute chain: decomposed solve converges and satisfies constraints.

        Revolute chains are under-constrained (each joint has 1 rotation DOF),
        so monolithic and decomposed may find different valid solutions.
        We verify convergence and that all constraint residuals are near zero.
        """
        pt_dec = ParamTable()
        bodies_dec, constraints_dec, indices_dec = self._build_chain(10, pt_dec)

        # Build clusters manually (chain -> each link is a biconnected component)
        cluster_list = []
        body_names = ["ground"] + [f"B{i}" for i in range(10)]
        for i in range(10):
            pair = {body_names[i], body_names[i + 1]}
            boundary = set()
            if i > 0:
                boundary.add(body_names[i])  # shared with previous cluster
            if i < 9:
                boundary.add(body_names[i + 1])  # shared with next cluster
            cluster_list.append(
                SolveCluster(
                    bodies=pair,
                    constraint_indices=[i],
                    boundary_bodies=boundary,
                    has_ground=(i == 0),
                )
            )

        converged = solve_decomposed(
            cluster_list,
            bodies_dec,
            constraints_dec,
            indices_dec,
            pt_dec,
        )
        assert converged

        # Verify all constraint residuals are satisfied
        env = pt_dec.get_env()
        for c in constraints_dec:
            for r in c.residuals():
                val = r.eval(env)
                assert abs(val) < 1e-8, f"Residual not satisfied: {val}"

        # Verify quat norms
        for name, body in bodies_dec.items():
            if not body.grounded:
                qn = body.quat_norm_residual().eval(env)
                assert abs(qn) < 1e-8, f"{name} quat not normalized: {qn}"

    def test_star_5_arms(self):
        """Star with 5 revolute arms off a grounded center."""
        pt = ParamTable()
        center = RigidBody("C", pt, (0, 0, 0), ID_QUAT, grounded=True)
        bodies = {"C": center}
        constraints = []
        indices = []
        clusters = []

        import math as m

        for i in range(5):
            name = f"arm{i}"
            angle = 2 * m.pi * i / 5
            marker_pos = (3 * m.cos(angle), 3 * m.sin(angle), 0)
            body = RigidBody(
                name, pt, (marker_pos[0] + 1, marker_pos[1] + 1, 0), ID_QUAT
            )
            bodies[name] = body

            c = RevoluteConstraint(
                center,
                marker_pos,
                ID_QUAT,
                body,
                (0, 0, 0),
                ID_QUAT,
            )
            constraints.append(c)
            indices.append(i)

            clusters.append(
                SolveCluster(
                    bodies={"C", name},
                    constraint_indices=[i],
                    boundary_bodies={"C"},
                    has_ground=True,
                )
            )

        converged = solve_decomposed(clusters, bodies, constraints, indices, pt)
        assert converged

        env = pt.get_env()
        for i in range(5):
            name = f"arm{i}"
            angle = 2 * m.pi * i / 5
            expected_x = 3 * m.cos(angle)
            expected_y = 3 * m.sin(angle)
            pos = bodies[name].extract_position(env)
            assert abs(pos[0] - expected_x) < 1e-6, (
                f"{name}: x={pos[0]:.4f} expected={expected_x:.4f}"
            )
            assert abs(pos[1] - expected_y) < 1e-6, (
                f"{name}: y={pos[1]:.4f} expected={expected_y:.4f}"
            )
            assert abs(pos[2]) < 1e-6

    def test_single_cluster_loop(self):
        """A loop assembly (single biconnected component) solves as one cluster."""
        pt = ParamTable()
        ground = RigidBody("g", pt, (0, 0, 0), ID_QUAT, grounded=True)
        b1 = RigidBody("B1", pt, (2.5, 0.3, 0), ID_QUAT)
        b2 = RigidBody("B2", pt, (4.5, -0.2, 0), ID_QUAT)
        bodies = {"g": ground, "B1": b1, "B2": b2}

        c0 = RevoluteConstraint(ground, (2, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT)
        c1 = RevoluteConstraint(b1, (2, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        c2 = RevoluteConstraint(b2, (2, 0, 0), ID_QUAT, ground, (6, 0, 0), ID_QUAT)
        constraint_objs = [c0, c1, c2]
        constraint_indices = [0, 1, 2]

        clusters = [
            SolveCluster(
                bodies={"g", "B1", "B2"},
                constraint_indices=[0, 1, 2],
                boundary_bodies=set(),
                has_ground=True,
            ),
        ]

        converged = solve_decomposed(
            clusters, bodies, constraint_objs, constraint_indices, pt
        )
        assert converged

        env = pt.get_env()
        for c in constraint_objs:
            for r in c.residuals():
                assert abs(r.eval(env)) < 1e-8

    def test_boundary_propagation_accuracy(self):
        """Verify boundary body values propagate correctly between clusters.

        Chain: ground → A → B → C. A is boundary between clusters 0-1.
        B is boundary between clusters 1-2. After decomposed solve,
        all revolute constraints should be satisfied.
        """
        pt = ParamTable()
        ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True)
        body_a = RigidBody("A", pt, (2.5, 0.5, 0), ID_QUAT)  # displaced
        body_b = RigidBody("B", pt, (4.5, -0.5, 0), ID_QUAT)  # displaced
        body_c = RigidBody("C", pt, (7.0, 1.0, 0), ID_QUAT)  # displaced
        bodies = {"ground": ground, "A": body_a, "B": body_b, "C": body_c}

        c0 = RevoluteConstraint(ground, (2, 0, 0), ID_QUAT, body_a, (0, 0, 0), ID_QUAT)
        c1 = RevoluteConstraint(body_a, (2, 0, 0), ID_QUAT, body_b, (0, 0, 0), ID_QUAT)
        c2 = RevoluteConstraint(body_b, (2, 0, 0), ID_QUAT, body_c, (0, 0, 0), ID_QUAT)
        constraint_objs = [c0, c1, c2]
        constraint_indices = [0, 1, 2]

        clusters = [
            SolveCluster(
                bodies={"ground", "A"},
                constraint_indices=[0],
                boundary_bodies={"A"},
                has_ground=True,
            ),
            SolveCluster(
                bodies={"A", "B"},
                constraint_indices=[1],
                boundary_bodies={"A", "B"},
                has_ground=False,
            ),
            SolveCluster(
                bodies={"B", "C"},
                constraint_indices=[2],
                boundary_bodies={"B"},
                has_ground=False,
            ),
        ]

        converged = solve_decomposed(
            clusters, bodies, constraint_objs, constraint_indices, pt
        )
        assert converged

        env = pt.get_env()
        # A at (2,0,0), B at (4,0,0), C at (6,0,0)
        pos_a = body_a.extract_position(env)
        pos_b = body_b.extract_position(env)
        pos_c = body_c.extract_position(env)

        assert abs(pos_a[0] - 2.0) < 1e-6
        assert abs(pos_a[1]) < 1e-6
        assert abs(pos_b[0] - 4.0) < 1e-6
        assert abs(pos_b[1]) < 1e-6
        assert abs(pos_c[0] - 6.0) < 1e-6
        assert abs(pos_c[1]) < 1e-6

        # Verify all params are free after solve (unfixed correctly)
        for name in ["A", "B", "C"]:
            body = bodies[name]
            for pname in body._param_names:
                assert not pt.is_fixed(pname), f"{pname} still fixed after solve"


# ============================================================================
# Pebble game rigidity classification
# ============================================================================


class TestPebbleGameClassification:
    """Test classify_cluster_rigidity using the pebble game."""

    def test_well_constrained_fixed_pair(self):
        """Two bodies with a Fixed joint (6 residuals) → well-constrained."""
        G = build_constraint_graph_simple(
            [("ground", "B", "Fixed", 0)],
            grounded={"ground"},
        )
        cluster = SolveCluster(
            bodies={"ground", "B"},
            constraint_indices=[0],
            boundary_bodies=set(),
            has_ground=True,
        )
        result = classify_cluster_rigidity(cluster, G, {"ground"})
        assert result == "well-constrained", f"Expected well-constrained, got {result}"

    def test_underconstrained_revolute_pair(self):
        """Ground + body with Revolute (5 residuals) → 1 DOF under-constrained."""
        G = build_constraint_graph_simple(
            [("ground", "B", "Revolute", 0)],
            grounded={"ground"},
        )
        cluster = SolveCluster(
            bodies={"ground", "B"},
            constraint_indices=[0],
            boundary_bodies=set(),
            has_ground=True,
        )
        result = classify_cluster_rigidity(cluster, G, {"ground"})
        assert result == "underconstrained", f"Expected underconstrained, got {result}"

    def test_overconstrained(self):
        """Two Fixed joints between same pair → overconstrained."""
        G = build_constraint_graph_simple(
            [
                ("ground", "B", "Fixed", 0),
                ("ground", "B", "Fixed", 1),
            ],
            grounded={"ground"},
        )
        cluster = SolveCluster(
            bodies={"ground", "B"},
            constraint_indices=[0, 1],
            boundary_bodies=set(),
            has_ground=True,
        )
        result = classify_cluster_rigidity(cluster, G, {"ground"})
        assert result == "overconstrained", f"Expected overconstrained, got {result}"

    def test_manual_edge_types(self):
        """Parallel (2 residuals, manual edge) is correctly handled."""
        G = build_constraint_graph_simple(
            [
                ("ground", "B", "Coincident", 0),  # 3 residuals
                ("ground", "B", "Parallel", 1),  # 2 residuals (manual)
            ],
            grounded={"ground"},
        )
        # Total: 3 + 2 = 5 residuals (like Revolute) → underconstrained (1 DOF)
        cluster = SolveCluster(
            bodies={"ground", "B"},
            constraint_indices=[0, 1],
            boundary_bodies=set(),
            has_ground=True,
        )
        result = classify_cluster_rigidity(cluster, G, {"ground"})
        assert result == "underconstrained", f"Expected underconstrained, got {result}"

    def test_no_ground(self):
        """Two ungrounded bodies with Fixed joint → well-constrained (6 DOF trivial)."""
        G = build_constraint_graph_simple(
            [("A", "B", "Fixed", 0)],
        )
        cluster = SolveCluster(
            bodies={"A", "B"},
            constraint_indices=[0],
            boundary_bodies=set(),
            has_ground=False,
        )
        result = classify_cluster_rigidity(cluster, G, set())
        # Two bodies, no ground: 12 DOF total, 6 from Fixed, 6 trivial → well-constrained
        assert result == "well-constrained", f"Expected well-constrained, got {result}"


# ============================================================================
# Large assembly tests (20+ bodies)
# ============================================================================


class TestLargeAssembly:
    """End-to-end tests with large synthetic assemblies."""

    def test_chain_20_bodies(self):
        """20-body revolute chain: decomposed solve converges."""
        pt = ParamTable()
        ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True)
        bodies = {"ground": ground}
        constraints = []
        indices = []

        prev = ground
        for i in range(20):
            name = f"B{i}"
            body = RigidBody(name, pt, ((i + 1) * 2 + 0.3, 0.2, 0), ID_QUAT)
            bodies[name] = body
            c = RevoluteConstraint(prev, (2, 0, 0), ID_QUAT, body, (0, 0, 0), ID_QUAT)
            constraints.append(c)
            indices.append(i)
            prev = body

        # Build clusters: chain of 20 biconnected components
        body_names = ["ground"] + [f"B{i}" for i in range(20)]
        cluster_list = []
        for i in range(20):
            pair = {body_names[i], body_names[i + 1]}
            boundary = set()
            if i > 0:
                boundary.add(body_names[i])
            if i < 19:
                boundary.add(body_names[i + 1])
            cluster_list.append(
                SolveCluster(
                    bodies=pair,
                    constraint_indices=[i],
                    boundary_bodies=boundary,
                    has_ground=(i == 0),
                )
            )

        converged = solve_decomposed(cluster_list, bodies, constraints, indices, pt)
        assert converged

        env = pt.get_env()
        for c in constraints:
            for r in c.residuals():
                assert abs(r.eval(env)) < 1e-8

    def test_star_10_arms(self):
        """Star with 10 revolute arms: 10 independent clusters."""
        pt = ParamTable()
        center = RigidBody("C", pt, (0, 0, 0), ID_QUAT, grounded=True)
        bodies = {"C": center}
        constraints = []
        indices = []
        clusters = []

        for i in range(10):
            name = f"arm{i}"
            angle = 2 * math.pi * i / 10
            mx = 3 * math.cos(angle)
            my = 3 * math.sin(angle)
            body = RigidBody(name, pt, (mx + 0.5, my + 0.5, 0), ID_QUAT)
            bodies[name] = body
            c = RevoluteConstraint(
                center, (mx, my, 0), ID_QUAT, body, (0, 0, 0), ID_QUAT
            )
            constraints.append(c)
            indices.append(i)
            clusters.append(
                SolveCluster(
                    bodies={"C", name},
                    constraint_indices=[i],
                    boundary_bodies={"C"},
                    has_ground=True,
                )
            )

        converged = solve_decomposed(clusters, bodies, constraints, indices, pt)
        assert converged

        env = pt.get_env()
        for i in range(10):
            name = f"arm{i}"
            angle = 2 * math.pi * i / 10
            pos = bodies[name].extract_position(env)
            assert abs(pos[0] - 3 * math.cos(angle)) < 1e-6
            assert abs(pos[1] - 3 * math.sin(angle)) < 1e-6

    def test_tree_assembly(self):
        """Tree: ground → A → (B, C), A → D. Mixed branching.

        Topology: ground-A is root cluster, A is articulation point.
        Three leaf clusters: {A,B}, {A,C}, {A,D}.
        """
        pt = ParamTable()
        ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True)
        a = RigidBody("A", pt, (2.5, 0.3, 0), ID_QUAT)
        b = RigidBody("B", pt, (4.5, 2.5, 0), ID_QUAT)
        c = RigidBody("C", pt, (4.5, -2.5, 0), ID_QUAT)
        d = RigidBody("D", pt, (4.5, 0.3, 0), ID_QUAT)
        bodies = {"ground": ground, "A": a, "B": b, "C": c, "D": d}

        c0 = RevoluteConstraint(ground, (2, 0, 0), ID_QUAT, a, (0, 0, 0), ID_QUAT)
        c1 = RevoluteConstraint(a, (2, 2, 0), ID_QUAT, b, (0, 0, 0), ID_QUAT)
        c2 = RevoluteConstraint(a, (2, -2, 0), ID_QUAT, c, (0, 0, 0), ID_QUAT)
        c3 = RevoluteConstraint(a, (2, 0, 0), ID_QUAT, d, (0, 0, 0), ID_QUAT)
        constraint_objs = [c0, c1, c2, c3]
        constraint_indices = [0, 1, 2, 3]

        # Root cluster: {ground, A}, then 3 leaf clusters
        clusters = [
            SolveCluster(
                bodies={"ground", "A"},
                constraint_indices=[0],
                boundary_bodies={"A"},
                has_ground=True,
            ),
            SolveCluster(
                bodies={"A", "B"},
                constraint_indices=[1],
                boundary_bodies={"A"},
                has_ground=False,
            ),
            SolveCluster(
                bodies={"A", "C"},
                constraint_indices=[2],
                boundary_bodies={"A"},
                has_ground=False,
            ),
            SolveCluster(
                bodies={"A", "D"},
                constraint_indices=[3],
                boundary_bodies={"A"},
                has_ground=False,
            ),
        ]

        converged = solve_decomposed(
            clusters, bodies, constraint_objs, constraint_indices, pt
        )
        assert converged

        env = pt.get_env()
        # A should be at (2, 0, 0)
        pos_a = a.extract_position(env)
        assert abs(pos_a[0] - 2.0) < 1e-6
        assert abs(pos_a[1]) < 1e-6

        # B at A + marker (2,2,0) = (4,2,0)
        pos_b = b.extract_position(env)
        assert abs(pos_b[0] - 4.0) < 1e-6
        assert abs(pos_b[1] - 2.0) < 1e-6

        # C at A + marker (2,-2,0) = (4,-2,0)
        pos_c = c.extract_position(env)
        assert abs(pos_c[0] - 4.0) < 1e-6
        assert abs(pos_c[1] + 2.0) < 1e-6

        # D at A + marker (2,0,0) = (4,0,0)
        pos_d = d.extract_position(env)
        assert abs(pos_d[0] - 4.0) < 1e-6
        assert abs(pos_d[1]) < 1e-6

    def test_mixed_constraint_chain(self):
        """Chain with mixed constraint types: Fixed, Coincident, Cylindrical."""
        pt = ParamTable()
        ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True)
        b1 = RigidBody("B1", pt, (2.5, 0.3, 0), ID_QUAT)
        b2 = RigidBody("B2", pt, (5.5, -0.2, 0), ID_QUAT)
        b3 = RigidBody("B3", pt, (8.5, 0.4, 0), ID_QUAT)
        bodies = {"ground": ground, "B1": b1, "B2": b2, "B3": b3}

        c0 = FixedConstraint(ground, (2, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT)
        c1 = CoincidentConstraint(b1, (3, 0, 0), b2, (0, 0, 0))
        c2 = CylindricalConstraint(b2, (3, 0, 0), ID_QUAT, b3, (0, 0, 0), ID_QUAT)
        constraint_objs = [c0, c1, c2]
        constraint_indices = [0, 1, 2]

        clusters = [
            SolveCluster(
                bodies={"ground", "B1"},
                constraint_indices=[0],
                boundary_bodies={"B1"},
                has_ground=True,
            ),
            SolveCluster(
                bodies={"B1", "B2"},
                constraint_indices=[1],
                boundary_bodies={"B1", "B2"},
                has_ground=False,
            ),
            SolveCluster(
                bodies={"B2", "B3"},
                constraint_indices=[2],
                boundary_bodies={"B2"},
                has_ground=False,
            ),
        ]

        converged = solve_decomposed(
            clusters, bodies, constraint_objs, constraint_indices, pt
        )
        assert converged

        env = pt.get_env()
        for c in constraint_objs:
            for r in c.residuals():
                assert abs(r.eval(env)) < 1e-8
