"""Tests for solver.datagen.analysis -- combined analysis function."""

from __future__ import annotations

import numpy as np
import pytest

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

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _two_bodies() -> list[RigidBody]:
    return [
        RigidBody(0, position=np.array([0.0, 0.0, 0.0])),
        RigidBody(1, position=np.array([2.0, 0.0, 0.0])),
    ]


def _triangle_bodies() -> list[RigidBody]:
    return [
        RigidBody(0, position=np.array([0.0, 0.0, 0.0])),
        RigidBody(1, position=np.array([2.0, 0.0, 0.0])),
        RigidBody(2, position=np.array([1.0, 1.7, 0.0])),
    ]


# ---------------------------------------------------------------------------
# Scenario 1: Two bodies + revolute (underconstrained, 1 internal DOF)
# ---------------------------------------------------------------------------


class TestTwoBodiesRevolute:
    """Demo scenario 1: two bodies connected by a revolute joint."""

    @pytest.fixture()
    def result(self) -> ConstraintAnalysis:
        bodies = _two_bodies()
        joints = [
            Joint(
                0,
                body_a=0,
                body_b=1,
                joint_type=JointType.REVOLUTE,
                anchor_a=np.array([1.0, 0.0, 0.0]),
                anchor_b=np.array([1.0, 0.0, 0.0]),
                axis=np.array([0.0, 0.0, 1.0]),
            ),
        ]
        return analyze_assembly(bodies, joints, ground_body=0)

    def test_internal_dof(self, result: ConstraintAnalysis) -> None:
        assert result.jacobian_internal_dof == 1

    def test_not_rigid(self, result: ConstraintAnalysis) -> None:
        assert not result.is_rigid

    def test_classification(self, result: ConstraintAnalysis) -> None:
        assert result.combinatorial_classification == "underconstrained"

    def test_no_redundant(self, result: ConstraintAnalysis) -> None:
        assert result.combinatorial_redundant == 0


# ---------------------------------------------------------------------------
# Scenario 2: Two bodies + fixed (well-constrained, 0 internal DOF)
# ---------------------------------------------------------------------------


class TestTwoBodiesFixed:
    """Demo scenario 2: two bodies connected by a fixed joint."""

    @pytest.fixture()
    def result(self) -> ConstraintAnalysis:
        bodies = _two_bodies()
        joints = [
            Joint(
                0,
                body_a=0,
                body_b=1,
                joint_type=JointType.FIXED,
                anchor_a=np.array([1.0, 0.0, 0.0]),
                anchor_b=np.array([1.0, 0.0, 0.0]),
            ),
        ]
        return analyze_assembly(bodies, joints, ground_body=0)

    def test_internal_dof(self, result: ConstraintAnalysis) -> None:
        assert result.jacobian_internal_dof == 0

    def test_rigid(self, result: ConstraintAnalysis) -> None:
        assert result.is_rigid

    def test_minimally_rigid(self, result: ConstraintAnalysis) -> None:
        assert result.is_minimally_rigid

    def test_classification(self, result: ConstraintAnalysis) -> None:
        assert result.combinatorial_classification == "well-constrained"


# ---------------------------------------------------------------------------
# Scenario 3: Triangle with revolute joints (overconstrained)
# ---------------------------------------------------------------------------


class TestTriangleRevolute:
    """Demo scenario 3: triangle of 3 bodies + 3 revolute joints."""

    @pytest.fixture()
    def result(self) -> ConstraintAnalysis:
        bodies = _triangle_bodies()
        joints = [
            Joint(
                0,
                body_a=0,
                body_b=1,
                joint_type=JointType.REVOLUTE,
                anchor_a=np.array([1.0, 0.0, 0.0]),
                anchor_b=np.array([1.0, 0.0, 0.0]),
                axis=np.array([0.0, 0.0, 1.0]),
            ),
            Joint(
                1,
                body_a=1,
                body_b=2,
                joint_type=JointType.REVOLUTE,
                anchor_a=np.array([1.5, 0.85, 0.0]),
                anchor_b=np.array([1.5, 0.85, 0.0]),
                axis=np.array([0.0, 0.0, 1.0]),
            ),
            Joint(
                2,
                body_a=2,
                body_b=0,
                joint_type=JointType.REVOLUTE,
                anchor_a=np.array([0.5, 0.85, 0.0]),
                anchor_b=np.array([0.5, 0.85, 0.0]),
                axis=np.array([0.0, 0.0, 1.0]),
            ),
        ]
        return analyze_assembly(bodies, joints, ground_body=0)

    def test_has_redundant(self, result: ConstraintAnalysis) -> None:
        assert result.combinatorial_redundant > 0

    def test_classification(self, result: ConstraintAnalysis) -> None:
        assert result.combinatorial_classification in ("overconstrained", "mixed")

    def test_rigid(self, result: ConstraintAnalysis) -> None:
        assert result.is_rigid

    def test_numerically_dependent(self, result: ConstraintAnalysis) -> None:
        assert len(result.numerically_dependent) > 0


# ---------------------------------------------------------------------------
# Scenario 4: Parallel revolute axes (geometric degeneracy)
# ---------------------------------------------------------------------------


class TestParallelRevoluteAxes:
    """Demo scenario 4: parallel revolute axes create geometric degeneracies."""

    @pytest.fixture()
    def result(self) -> ConstraintAnalysis:
        bodies = [
            RigidBody(0, position=np.array([0.0, 0.0, 0.0])),
            RigidBody(1, position=np.array([2.0, 0.0, 0.0])),
            RigidBody(2, position=np.array([4.0, 0.0, 0.0])),
        ]
        joints = [
            Joint(
                0,
                body_a=0,
                body_b=1,
                joint_type=JointType.REVOLUTE,
                anchor_a=np.array([1.0, 0.0, 0.0]),
                anchor_b=np.array([1.0, 0.0, 0.0]),
                axis=np.array([0.0, 0.0, 1.0]),
            ),
            Joint(
                1,
                body_a=1,
                body_b=2,
                joint_type=JointType.REVOLUTE,
                anchor_a=np.array([3.0, 0.0, 0.0]),
                anchor_b=np.array([3.0, 0.0, 0.0]),
                axis=np.array([0.0, 0.0, 1.0]),
            ),
        ]
        return analyze_assembly(bodies, joints, ground_body=0)

    def test_geometric_degeneracies_detected(self, result: ConstraintAnalysis) -> None:
        """Parallel axes produce at least one geometric degeneracy."""
        assert result.geometric_degeneracies > 0


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------


class TestNoJoints:
    """Assembly with bodies but no joints."""

    def test_all_dof_free(self) -> None:
        bodies = _two_bodies()
        result = analyze_assembly(bodies, [], ground_body=0)
        # Body 1 is completely free (6 DOF), body 0 is grounded
        assert result.jacobian_internal_dof > 0
        assert not result.is_rigid

    def test_ungrounded(self) -> None:
        bodies = _two_bodies()
        result = analyze_assembly(bodies, [])
        assert result.combinatorial_classification == "underconstrained"


class TestReturnType:
    """Verify the return object is a proper ConstraintAnalysis."""

    def test_instance(self) -> None:
        bodies = _two_bodies()
        joints = [Joint(0, 0, 1, JointType.FIXED)]
        result = analyze_assembly(bodies, joints)
        assert isinstance(result, ConstraintAnalysis)

    def test_per_edge_results_populated(self) -> None:
        bodies = _two_bodies()
        joints = [Joint(0, 0, 1, JointType.REVOLUTE)]
        result = analyze_assembly(bodies, joints)
        assert len(result.per_edge_results) == 5
