"""Tests for solver.datagen.pebble_game -- (6,6)-pebble game."""

from __future__ import annotations

import itertools

import numpy as np
import pytest

from solver.datagen.pebble_game import PebbleGame3D
from solver.datagen.types import Joint, JointType

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


def _revolute(jid: int, a: int, b: int, axis: np.ndarray | None = None) -> Joint:
    """Shorthand for a revolute joint between bodies *a* and *b*."""
    if axis is None:
        axis = np.array([0.0, 0.0, 1.0])
    return Joint(
        joint_id=jid,
        body_a=a,
        body_b=b,
        joint_type=JointType.REVOLUTE,
        axis=axis,
    )


def _fixed(jid: int, a: int, b: int) -> Joint:
    return Joint(joint_id=jid, body_a=a, body_b=b, joint_type=JointType.FIXED)


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestAddBody:
    """Body registration basics."""

    def test_single_body_six_pebbles(self) -> None:
        pg = PebbleGame3D()
        pg.add_body(0)
        assert pg.state.free_pebbles[0] == 6

    def test_duplicate_body_no_op(self) -> None:
        pg = PebbleGame3D()
        pg.add_body(0)
        pg.add_body(0)
        assert pg.state.free_pebbles[0] == 6

    def test_multiple_bodies(self) -> None:
        pg = PebbleGame3D()
        for i in range(5):
            pg.add_body(i)
        assert pg.get_dof() == 30  # 5 * 6


class TestAddJoint:
    """Joint insertion and DOF accounting."""

    def test_revolute_removes_five_dof(self) -> None:
        pg = PebbleGame3D()
        results = pg.add_joint(_revolute(0, 0, 1))
        assert len(results) == 5  # 5 scalar constraints
        assert all(r["independent"] for r in results)
        # 2 bodies * 6 = 12, minus 5 independent = 7 free pebbles
        assert pg.get_dof() == 7

    def test_fixed_removes_six_dof(self) -> None:
        pg = PebbleGame3D()
        results = pg.add_joint(_fixed(0, 0, 1))
        assert len(results) == 6
        assert all(r["independent"] for r in results)
        assert pg.get_dof() == 6

    def test_ball_removes_three_dof(self) -> None:
        pg = PebbleGame3D()
        j = Joint(joint_id=0, body_a=0, body_b=1, joint_type=JointType.BALL)
        results = pg.add_joint(j)
        assert len(results) == 3
        assert all(r["independent"] for r in results)
        assert pg.get_dof() == 9


class TestTwoBodiesRevolute:
    """Two bodies connected by a revolute -- demo scenario 1."""

    def test_internal_dof(self) -> None:
        pg = PebbleGame3D()
        pg.add_joint(_revolute(0, 0, 1))
        # Total DOF = 7, internal = 7 - 6 = 1
        assert pg.get_internal_dof() == 1

    def test_not_rigid(self) -> None:
        pg = PebbleGame3D()
        pg.add_joint(_revolute(0, 0, 1))
        assert not pg.is_rigid()

    def test_classification(self) -> None:
        pg = PebbleGame3D()
        pg.add_joint(_revolute(0, 0, 1))
        assert pg.classify_assembly() == "underconstrained"


class TestTwoBodiesFixed:
    """Two bodies + fixed joint -- demo scenario 2."""

    def test_zero_internal_dof(self) -> None:
        pg = PebbleGame3D()
        pg.add_joint(_fixed(0, 0, 1))
        assert pg.get_internal_dof() == 0

    def test_rigid(self) -> None:
        pg = PebbleGame3D()
        pg.add_joint(_fixed(0, 0, 1))
        assert pg.is_rigid()

    def test_well_constrained(self) -> None:
        pg = PebbleGame3D()
        pg.add_joint(_fixed(0, 0, 1))
        assert pg.classify_assembly() == "well-constrained"


class TestTriangleRevolute:
    """Triangle of 3 bodies with revolute joints -- demo scenario 3."""

    @pytest.fixture()
    def pg(self) -> PebbleGame3D:
        pg = PebbleGame3D()
        pg.add_joint(_revolute(0, 0, 1))
        pg.add_joint(_revolute(1, 1, 2))
        pg.add_joint(_revolute(2, 2, 0))
        return pg

    def test_has_redundant_edges(self, pg: PebbleGame3D) -> None:
        assert pg.get_redundant_count() > 0

    def test_classification_overconstrained(self, pg: PebbleGame3D) -> None:
        # 15 constraints on 3 bodies (Maxwell: 6*3-6=12 needed)
        assert pg.classify_assembly() in ("overconstrained", "mixed")

    def test_rigid(self, pg: PebbleGame3D) -> None:
        assert pg.is_rigid()


class TestChainNotRigid:
    """A serial chain of 4 bodies with revolute joints is never rigid."""

    def test_chain_underconstrained(self) -> None:
        pg = PebbleGame3D()
        for i in range(3):
            pg.add_joint(_revolute(i, i, i + 1))
        assert not pg.is_rigid()
        assert pg.classify_assembly() == "underconstrained"

    def test_chain_internal_dof(self) -> None:
        pg = PebbleGame3D()
        for i in range(3):
            pg.add_joint(_revolute(i, i, i + 1))
        # 4 bodies * 6 = 24, minus 15 independent = 9 free, internal = 3
        assert pg.get_internal_dof() == 3


class TestEdgeResults:
    """Result dicts returned by add_joint."""

    def test_result_keys(self) -> None:
        pg = PebbleGame3D()
        results = pg.add_joint(_revolute(0, 0, 1))
        expected_keys = {"edge_id", "joint_id", "constraint_index", "independent", "dof_remaining"}
        for r in results:
            assert set(r.keys()) == expected_keys

    def test_edge_ids_sequential(self) -> None:
        pg = PebbleGame3D()
        r1 = pg.add_joint(_revolute(0, 0, 1))
        r2 = pg.add_joint(_revolute(1, 1, 2))
        all_ids = [r["edge_id"] for r in r1 + r2]
        assert all_ids == list(range(10))

    def test_dof_remaining_monotonic(self) -> None:
        pg = PebbleGame3D()
        results = pg.add_joint(_revolute(0, 0, 1))
        dofs = [r["dof_remaining"] for r in results]
        # Should be non-increasing (each independent edge removes a pebble)
        for a, b in itertools.pairwise(dofs):
            assert a >= b


class TestGroundedClassification:
    """classify_assembly with grounded=True."""

    def test_grounded_baseline_zero(self) -> None:
        """With grounded=True the baseline is 0 (not 6)."""
        pg = PebbleGame3D()
        pg.add_joint(_fixed(0, 0, 1))
        # Ungrounded: well-constrained (6 pebbles = baseline 6)
        assert pg.classify_assembly(grounded=False) == "well-constrained"
        # Grounded: the 6 remaining pebbles on body 1 exceed baseline 0,
        # so the raw pebble game (without a virtual ground body) sees this
        # as underconstrained. The analysis function handles this properly
        # by adding a virtual ground body.
        assert pg.classify_assembly(grounded=True) == "underconstrained"
