"""Tests for solver.datagen.types -- shared data types."""

from __future__ import annotations

from typing import ClassVar

import numpy as np
import pytest

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


class TestJointType:
    """JointType enum construction and DOF values."""

    EXPECTED_DOF: ClassVar[dict[str, int]] = {
        "FIXED": 6,
        "REVOLUTE": 5,
        "CYLINDRICAL": 4,
        "SLIDER": 5,
        "BALL": 3,
        "PLANAR": 3,
        "SCREW": 5,
        "UNIVERSAL": 4,
        "PARALLEL": 3,
        "PERPENDICULAR": 1,
        "DISTANCE": 1,
    }

    def test_member_count(self) -> None:
        assert len(JointType) == 11

    @pytest.mark.parametrize("name,dof", EXPECTED_DOF.items())
    def test_dof_values(self, name: str, dof: int) -> None:
        assert JointType[name].dof == dof

    def test_access_by_name(self) -> None:
        assert JointType["REVOLUTE"] is JointType.REVOLUTE

    def test_value_is_tuple(self) -> None:
        assert JointType.REVOLUTE.value == (1, 5)
        assert JointType.REVOLUTE.dof == 5


class TestRigidBody:
    """RigidBody dataclass defaults and construction."""

    def test_defaults(self) -> None:
        body = RigidBody(body_id=0)
        np.testing.assert_array_equal(body.position, np.zeros(3))
        np.testing.assert_array_equal(body.orientation, np.eye(3))
        assert body.local_anchors == {}

    def test_custom_position(self) -> None:
        pos = np.array([1.0, 2.0, 3.0])
        body = RigidBody(body_id=7, position=pos)
        np.testing.assert_array_equal(body.position, pos)
        assert body.body_id == 7

    def test_local_anchors_mutable(self) -> None:
        body = RigidBody(body_id=0)
        body.local_anchors["top"] = np.array([0.0, 0.0, 1.0])
        assert "top" in body.local_anchors

    def test_default_factory_isolation(self) -> None:
        """Each instance gets its own default containers."""
        b1 = RigidBody(body_id=0)
        b2 = RigidBody(body_id=1)
        b1.local_anchors["x"] = np.zeros(3)
        assert "x" not in b2.local_anchors


class TestJoint:
    """Joint dataclass defaults and construction."""

    def test_defaults(self) -> None:
        j = Joint(joint_id=0, body_a=0, body_b=1, joint_type=JointType.REVOLUTE)
        np.testing.assert_array_equal(j.anchor_a, np.zeros(3))
        np.testing.assert_array_equal(j.anchor_b, np.zeros(3))
        np.testing.assert_array_equal(j.axis, np.array([0.0, 0.0, 1.0]))
        assert j.pitch == 0.0

    def test_full_construction(self) -> None:
        j = Joint(
            joint_id=5,
            body_a=2,
            body_b=3,
            joint_type=JointType.SCREW,
            anchor_a=np.array([1.0, 0.0, 0.0]),
            anchor_b=np.array([2.0, 0.0, 0.0]),
            axis=np.array([1.0, 0.0, 0.0]),
            pitch=0.5,
        )
        assert j.joint_id == 5
        assert j.joint_type is JointType.SCREW
        assert j.pitch == 0.5


class TestPebbleState:
    """PebbleState dataclass defaults."""

    def test_defaults(self) -> None:
        s = PebbleState()
        assert s.free_pebbles == {}
        assert s.directed_edges == {}
        assert s.independent_edges == set()
        assert s.redundant_edges == set()
        assert s.incoming == {}
        assert s.outgoing == {}

    def test_default_factory_isolation(self) -> None:
        s1 = PebbleState()
        s2 = PebbleState()
        s1.free_pebbles[0] = 6
        assert 0 not in s2.free_pebbles


class TestConstraintAnalysis:
    """ConstraintAnalysis dataclass construction."""

    def test_construction(self) -> None:
        ca = ConstraintAnalysis(
            combinatorial_dof=6,
            combinatorial_internal_dof=0,
            combinatorial_redundant=0,
            combinatorial_classification="well-constrained",
            per_edge_results=[],
            jacobian_rank=6,
            jacobian_nullity=0,
            jacobian_internal_dof=0,
            numerically_dependent=[],
            geometric_degeneracies=0,
            is_rigid=True,
            is_minimally_rigid=True,
        )
        assert ca.is_rigid is True
        assert ca.is_minimally_rigid is True
        assert ca.combinatorial_classification == "well-constrained"

    def test_per_edge_results_typing(self) -> None:
        """per_edge_results accepts list[dict[str, Any]]."""
        ca = ConstraintAnalysis(
            combinatorial_dof=7,
            combinatorial_internal_dof=1,
            combinatorial_redundant=0,
            combinatorial_classification="underconstrained",
            per_edge_results=[{"edge_id": 0, "independent": True}],
            jacobian_rank=5,
            jacobian_nullity=1,
            jacobian_internal_dof=1,
            numerically_dependent=[],
            geometric_degeneracies=0,
            is_rigid=False,
            is_minimally_rigid=False,
        )
        assert len(ca.per_edge_results) == 1
        assert ca.per_edge_results[0]["edge_id"] == 0
