"""Tests for per-entity DOF diagnostics and overconstrained detection."""

import math

import numpy as np
import pytest
from kindred_solver.constraints import (
    CoincidentConstraint,
    CylindricalConstraint,
    DistancePointPointConstraint,
    FixedConstraint,
    ParallelConstraint,
    RevoluteConstraint,
)
from kindred_solver.diagnostics import (
    ConstraintDiag,
    EntityDOF,
    find_overconstrained,
    per_entity_dof,
)
from kindred_solver.entities import RigidBody
from kindred_solver.params import ParamTable


def _make_two_bodies(
    params,
    pos_a=(0, 0, 0),
    pos_b=(5, 0, 0),
    quat_a=(1, 0, 0, 0),
    quat_b=(1, 0, 0, 0),
    ground_a=True,
    ground_b=False,
):
    body_a = RigidBody(
        "a", params, position=pos_a, quaternion=quat_a, grounded=ground_a
    )
    body_b = RigidBody(
        "b", params, position=pos_b, quaternion=quat_b, grounded=ground_b
    )
    return body_a, body_b


def _build_residuals_and_ranges(constraint_objs, bodies, params):
    """Build residuals list, quat norms, and residual_ranges."""
    all_residuals = []
    residual_ranges = []
    row = 0
    for i, obj in enumerate(constraint_objs):
        r = obj.residuals()
        n = len(r)
        residual_ranges.append((row, row + n, i))
        all_residuals.extend(r)
        row += n

    for body in bodies.values():
        if not body.grounded:
            all_residuals.append(body.quat_norm_residual())

    return all_residuals, residual_ranges


# ============================================================================
# Per-entity DOF tests
# ============================================================================


class TestPerEntityDOF:
    """Per-entity DOF computation."""

    def test_unconstrained_body_6dof(self):
        """Unconstrained non-grounded body has 6 DOF."""
        params = ParamTable()
        body = RigidBody(
            "b", params, position=(0, 0, 0), quaternion=(1, 0, 0, 0), grounded=False
        )
        bodies = {"b": body}

        # Only quat norm constraint
        residuals = [body.quat_norm_residual()]

        result = per_entity_dof(residuals, params, bodies)
        assert len(result) == 1
        assert result[0].entity_id == "b"
        assert result[0].remaining_dof == 6
        assert len(result[0].free_motions) == 6

    def test_fixed_body_0dof(self):
        """Body welded to ground has 0 DOF."""
        params = ParamTable()
        body_a, body_b = _make_two_bodies(params)
        bodies = {"a": body_a, "b": body_b}

        c = FixedConstraint(
            body_a,
            (0, 0, 0),
            (1, 0, 0, 0),
            body_b,
            (0, 0, 0),
            (1, 0, 0, 0),
        )
        residuals, _ = _build_residuals_and_ranges([c], bodies, params)

        result = per_entity_dof(residuals, params, bodies)
        # Only non-grounded body (b) reported
        assert len(result) == 1
        assert result[0].entity_id == "b"
        assert result[0].remaining_dof == 0
        assert len(result[0].free_motions) == 0

    def test_revolute_1dof(self):
        """Revolute joint leaves 1 DOF (rotation about Z)."""
        params = ParamTable()
        body_a, body_b = _make_two_bodies(params, pos_b=(0, 0, 0))
        bodies = {"a": body_a, "b": body_b}

        c = RevoluteConstraint(
            body_a,
            (0, 0, 0),
            (1, 0, 0, 0),
            body_b,
            (0, 0, 0),
            (1, 0, 0, 0),
        )
        residuals, _ = _build_residuals_and_ranges([c], bodies, params)

        result = per_entity_dof(residuals, params, bodies)
        assert len(result) == 1
        assert result[0].remaining_dof == 1
        # Should have one free motion that mentions rotation
        assert len(result[0].free_motions) == 1
        assert "rotation" in result[0].free_motions[0].lower()

    def test_cylindrical_2dof(self):
        """Cylindrical joint leaves 2 DOF (rotation about Z + translation along Z)."""
        params = ParamTable()
        body_a, body_b = _make_two_bodies(params, pos_b=(0, 0, 0))
        bodies = {"a": body_a, "b": body_b}

        c = CylindricalConstraint(
            body_a,
            (0, 0, 0),
            (1, 0, 0, 0),
            body_b,
            (0, 0, 0),
            (1, 0, 0, 0),
        )
        residuals, _ = _build_residuals_and_ranges([c], bodies, params)

        result = per_entity_dof(residuals, params, bodies)
        assert len(result) == 1
        assert result[0].remaining_dof == 2
        assert len(result[0].free_motions) == 2

    def test_coincident_3dof(self):
        """Coincident (ball) joint leaves 3 DOF (3 rotations)."""
        params = ParamTable()
        body_a, body_b = _make_two_bodies(params, pos_b=(0, 0, 0))
        bodies = {"a": body_a, "b": body_b}

        c = CoincidentConstraint(body_a, (0, 0, 0), body_b, (0, 0, 0))
        residuals, _ = _build_residuals_and_ranges([c], bodies, params)

        result = per_entity_dof(residuals, params, bodies)
        assert len(result) == 1
        assert result[0].remaining_dof == 3
        # All 3 should be rotations
        for motion in result[0].free_motions:
            assert "rotation" in motion.lower()

    def test_no_constraints_6dof(self):
        """No residuals at all gives 6 DOF."""
        params = ParamTable()
        body = RigidBody(
            "b", params, position=(0, 0, 0), quaternion=(1, 0, 0, 0), grounded=False
        )
        bodies = {"b": body}

        result = per_entity_dof([], params, bodies)
        assert len(result) == 1
        assert result[0].remaining_dof == 6

    def test_grounded_body_excluded(self):
        """Grounded bodies are not reported."""
        params = ParamTable()
        body_a, body_b = _make_two_bodies(params)
        bodies = {"a": body_a, "b": body_b}

        residuals = [body_b.quat_norm_residual()]
        result = per_entity_dof(residuals, params, bodies)

        entity_ids = [r.entity_id for r in result]
        assert "a" not in entity_ids  # grounded
        assert "b" in entity_ids

    def test_multiple_bodies(self):
        """Two free bodies: each gets its own DOF report."""
        params = ParamTable()
        body_g = RigidBody(
            "g", params, position=(0, 0, 0), quaternion=(1, 0, 0, 0), grounded=True
        )
        body_b = RigidBody(
            "b", params, position=(5, 0, 0), quaternion=(1, 0, 0, 0), grounded=False
        )
        body_c = RigidBody(
            "c", params, position=(10, 0, 0), quaternion=(1, 0, 0, 0), grounded=False
        )
        bodies = {"g": body_g, "b": body_b, "c": body_c}

        # Fix b to ground, leave c unconstrained
        c_fix = FixedConstraint(
            body_g,
            (0, 0, 0),
            (1, 0, 0, 0),
            body_b,
            (0, 0, 0),
            (1, 0, 0, 0),
        )
        residuals, _ = _build_residuals_and_ranges([c_fix], bodies, params)

        result = per_entity_dof(residuals, params, bodies)
        result_map = {r.entity_id: r for r in result}

        assert result_map["b"].remaining_dof == 0
        assert result_map["c"].remaining_dof == 6


# ============================================================================
# Overconstrained detection tests
# ============================================================================


class TestFindOverconstrained:
    """Redundant and conflicting constraint detection."""

    def test_well_constrained_no_diagnostics(self):
        """Well-constrained system produces no diagnostics."""
        params = ParamTable()
        body_a, body_b = _make_two_bodies(params, pos_b=(0, 0, 0))
        bodies = {"a": body_a, "b": body_b}

        c = FixedConstraint(
            body_a,
            (0, 0, 0),
            (1, 0, 0, 0),
            body_b,
            (0, 0, 0),
            (1, 0, 0, 0),
        )
        residuals, ranges = _build_residuals_and_ranges([c], bodies, params)

        diags = find_overconstrained(residuals, params, ranges)
        assert len(diags) == 0

    def test_duplicate_coincident_redundant(self):
        """Duplicate coincident constraint is flagged as redundant."""
        params = ParamTable()
        body_a, body_b = _make_two_bodies(params, pos_b=(0, 0, 0))
        bodies = {"a": body_a, "b": body_b}

        c1 = CoincidentConstraint(body_a, (0, 0, 0), body_b, (0, 0, 0))
        c2 = CoincidentConstraint(body_a, (0, 0, 0), body_b, (0, 0, 0))
        residuals, ranges = _build_residuals_and_ranges([c1, c2], bodies, params)

        diags = find_overconstrained(residuals, params, ranges)
        assert len(diags) > 0
        # At least one should be redundant
        kinds = {d.kind for d in diags}
        assert "redundant" in kinds

    def test_conflicting_distance(self):
        """Distance constraint that can't be satisfied is flagged as conflicting."""
        params = ParamTable()
        body_a, body_b = _make_two_bodies(params, pos_b=(0, 0, 0))
        bodies = {"a": body_a, "b": body_b}

        # Coincident forces distance=0, but distance constraint says 50
        c1 = CoincidentConstraint(body_a, (0, 0, 0), body_b, (0, 0, 0))
        c2 = DistancePointPointConstraint(
            body_a,
            (0, 0, 0),
            body_b,
            (0, 0, 0),
            distance=50.0,
        )
        residuals, ranges = _build_residuals_and_ranges([c1, c2], bodies, params)

        diags = find_overconstrained(residuals, params, ranges)
        assert len(diags) > 0
        kinds = {d.kind for d in diags}
        assert "conflicting" in kinds

    def test_empty_system_no_diagnostics(self):
        """Empty system has no diagnostics."""
        params = ParamTable()
        diags = find_overconstrained([], params, [])
        assert len(diags) == 0