"""Tests for Phase 2 constraint residual generation."""

import math

import pytest
from kindred_solver.constraints import (
    AngleConstraint,
    BallConstraint,
    CamConstraint,
    ConcentricConstraint,
    CylindricalConstraint,
    DistanceCylSphConstraint,
    GearConstraint,
    LineInPlaneConstraint,
    ParallelConstraint,
    PerpendicularConstraint,
    PlanarConstraint,
    PointInPlaneConstraint,
    PointOnLineConstraint,
    RackPinionConstraint,
    RevoluteConstraint,
    ScrewConstraint,
    SliderConstraint,
    SlotConstraint,
    TangentConstraint,
    UniversalConstraint,
)
from kindred_solver.entities import RigidBody
from kindred_solver.params import ParamTable

ID_QUAT = (1.0, 0.0, 0.0, 0.0)
# 90-deg about Y: Z-axis of body rotates to point along X
_c = math.cos(math.pi / 4)
_s = math.sin(math.pi / 4)
ROT_90Y = (_c, 0.0, _s, 0.0)
ROT_90Z = (_c, 0.0, 0.0, _s)


# ── Point constraints ────────────────────────────────────────────────


class TestPointOnLine:
    def test_on_line(self):
        """Point at (0,0,5) is on Z-axis line through origin."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 5), (1, 0, 0, 0))
        c = PointOnLineConstraint(b2, (0, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_off_line(self):
        """Point at (3,0,5) is NOT on Z-axis line through origin."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (3, 0, 5), (1, 0, 0, 0))
        c = PointOnLineConstraint(b2, (0, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        vals = [r.eval(env) for r in c.residuals()]
        assert any(abs(v) > 0.1 for v in vals)

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = PointOnLineConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        assert len(c.residuals()) == 3


class TestPointInPlane:
    def test_in_plane(self):
        """Point at (3,4,0) is in XY plane through origin."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (3, 4, 0), (1, 0, 0, 0))
        c = PointInPlaneConstraint(b2, (0, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        assert abs(c.residuals()[0].eval(env)) < 1e-10

    def test_above_plane(self):
        """Point at (0,0,7) is 7 above XY plane."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 7), (1, 0, 0, 0))
        c = PointInPlaneConstraint(b2, (0, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        assert abs(c.residuals()[0].eval(env) - 7.0) < 1e-10

    def test_with_offset(self):
        """Point at (0,0,5) with offset=5 → residual 0."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 5), (1, 0, 0, 0))
        c = PointInPlaneConstraint(
            b2, (0, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT, offset=5.0
        )
        env = pt.get_env()
        assert abs(c.residuals()[0].eval(env)) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = PointInPlaneConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        assert len(c.residuals()) == 1


# ── Orientation constraints ──────────────────────────────────────────


class TestParallel:
    def test_parallel_same(self):
        """Both bodies with identity rotation → Z-axes parallel → residuals 0."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (5, 0, 0), (1, 0, 0, 0))
        c = ParallelConstraint(b1, ID_QUAT, b2, ID_QUAT)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_not_parallel(self):
        """One body rotated 90-deg about Y → Z-axes perpendicular."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (5, 0, 0), ROT_90Y)
        c = ParallelConstraint(b1, ID_QUAT, b2, ID_QUAT)
        env = pt.get_env()
        vals = [r.eval(env) for r in c.residuals()]
        assert any(abs(v) > 0.1 for v in vals)

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = ParallelConstraint(b1, ID_QUAT, b2, ID_QUAT)
        assert len(c.residuals()) == 3


class TestPerpendicular:
    def test_perpendicular(self):
        """One body rotated 90-deg about Y → Z-axes perpendicular."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), ROT_90Y)
        c = PerpendicularConstraint(b1, ID_QUAT, b2, ID_QUAT)
        env = pt.get_env()
        assert abs(c.residuals()[0].eval(env)) < 1e-10

    def test_not_perpendicular(self):
        """Same orientation → not perpendicular."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = PerpendicularConstraint(b1, ID_QUAT, b2, ID_QUAT)
        env = pt.get_env()
        # dot(z,z) = 1 ≠ 0
        assert abs(c.residuals()[0].eval(env) - 1.0) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = PerpendicularConstraint(b1, ID_QUAT, b2, ID_QUAT)
        assert len(c.residuals()) == 1


class TestAngle:
    def test_90_degrees(self):
        """90-deg angle between Z-axes rotated 90-deg about Y."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), ROT_90Y)
        c = AngleConstraint(b1, ID_QUAT, b2, ID_QUAT, math.pi / 2)
        env = pt.get_env()
        assert abs(c.residuals()[0].eval(env)) < 1e-10

    def test_0_degrees(self):
        """0-deg angle, same orientation → cos(0)=1, dot=1 → residual 0."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = AngleConstraint(b1, ID_QUAT, b2, ID_QUAT, 0.0)
        env = pt.get_env()
        assert abs(c.residuals()[0].eval(env)) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = AngleConstraint(b1, ID_QUAT, b2, ID_QUAT, 1.0)
        assert len(c.residuals()) == 1


# ── Axis/surface constraints ─────────────────────────────────────────


class TestConcentric:
    def test_coaxial(self):
        """Both on Z-axis → coaxial → residuals 0."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 5), (1, 0, 0, 0))
        c = ConcentricConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_not_coaxial(self):
        """Offset in X → not coaxial."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (5, 0, 0), (1, 0, 0, 0))
        c = ConcentricConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        vals = [r.eval(env) for r in c.residuals()]
        assert any(abs(v) > 0.1 for v in vals)

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = ConcentricConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        assert len(c.residuals()) == 6


class TestTangent:
    def test_touching(self):
        """Marker origins at same point → tangent."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = TangentConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        assert abs(c.residuals()[0].eval(env)) < 1e-10

    def test_separated(self):
        """Separated along normal → non-zero residual."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 5), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = TangentConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        assert abs(c.residuals()[0].eval(env) - 5.0) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = TangentConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        assert len(c.residuals()) == 1


class TestPlanar:
    def test_coplanar(self):
        """Same plane, same orientation → all residuals 0."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (5, 3, 0), (1, 0, 0, 0))
        c = PlanarConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_with_offset(self):
        """b_i at z=5, b_j at origin, normal=Z, offset=5.
        Signed distance = (p_i - p_j).n = 5, offset=5 → 5-5 = 0."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 5), (1, 0, 0, 0))
        c = PlanarConstraint(b2, (0, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT, offset=5.0)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = PlanarConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        assert len(c.residuals()) == 4


class TestLineInPlane:
    def test_in_plane(self):
        """Line along X in XY plane → residuals 0."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        # b2 has Z-axis = (1,0,0) via 90-deg rotation about Y
        b2 = RigidBody("b", pt, (5, 0, 0), ROT_90Y)
        # Line = b2's Z-axis (which is world X), plane = b1's XY plane (normal=Z)
        c = LineInPlaneConstraint(b2, (0, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = LineInPlaneConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        assert len(c.residuals()) == 2


# ── Kinematic joints ─────────────────────────────────────────────────


class TestBall:
    def test_same_as_coincident(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = BallConstraint(b1, (0, 0, 0), b2, (0, 0, 0))
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10
        assert len(c.residuals()) == 3


class TestRevolute:
    def test_satisfied(self):
        """Same position, same Z-axis → satisfied."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), ROT_90Z)  # rotated about Z — still parallel
        c = RevoluteConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = RevoluteConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        assert len(c.residuals()) == 6


class TestCylindrical:
    def test_on_axis(self):
        """Same axis, displaced along Z → satisfied."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 10), (1, 0, 0, 0))
        c = CylindricalConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = CylindricalConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        assert len(c.residuals()) == 6


class TestSlider:
    def test_aligned(self):
        """Same axis, no twist, displaced along Z → satisfied."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 10), (1, 0, 0, 0))
        c = SliderConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_twisted(self):
        """Rotated about Z → twist residual non-zero."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), ROT_90Z)
        c = SliderConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        vals = [r.eval(env) for r in c.residuals()]
        # First 6 should be ~0 (parallel + on-line), but twist residual should be ~1
        assert abs(vals[6]) > 0.5

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = SliderConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        assert len(c.residuals()) == 7


class TestUniversal:
    def test_satisfied(self):
        """Same origin, perpendicular Z-axes."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), ROT_90Y)
        c = UniversalConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = UniversalConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT)
        assert len(c.residuals()) == 4


class TestScrew:
    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = ScrewConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT, pitch=10.0)
        assert len(c.residuals()) == 7

    def test_zero_displacement_zero_rotation(self):
        """Both at origin with identity rotation → all residuals 0."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = ScrewConstraint(b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT, pitch=10.0)
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10


# ── Mechanical constraints ───────────────────────────────────────────


class TestGear:
    def test_both_at_rest(self):
        """Both at identity rotation → residual 0."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = GearConstraint(b1, ID_QUAT, b2, ID_QUAT, 1.0, 1.0)
        env = pt.get_env()
        assert abs(c.residuals()[0].eval(env)) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = GearConstraint(b1, ID_QUAT, b2, ID_QUAT, 1.0, 2.0)
        assert len(c.residuals()) == 1


class TestRackPinion:
    def test_at_rest(self):
        """Both at rest → residual 0."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = RackPinionConstraint(
            b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT, pitch_radius=5.0
        )
        env = pt.get_env()
        assert abs(c.residuals()[0].eval(env)) < 1e-10

    def test_residual_count(self):
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        b2 = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = RackPinionConstraint(
            b1, (0, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT, pitch_radius=1.0
        )
        assert len(c.residuals()) == 1


# ── Stubs ────────────────────────────────────────────────────────────


class TestStubs:
    def test_cam(self):
        assert CamConstraint().residuals() == []

    def test_slot(self):
        assert SlotConstraint().residuals() == []

    def test_distance_cyl_sph(self):
        assert DistanceCylSphConstraint().residuals() == []