"""Tests for constraint residual generation."""

import math

import pytest
from kindred_solver.constraints import (
    CoincidentConstraint,
    DistancePointPointConstraint,
    FixedConstraint,
)
from kindred_solver.entities import RigidBody
from kindred_solver.params import ParamTable


class TestCoincident:
    def test_satisfied(self):
        """Two bodies at same position → residuals are zero."""
        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 = CoincidentConstraint(b1, (0, 0, 0), b2, (0, 0, 0))
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_unsatisfied(self):
        """Bodies at different positions → non-zero residuals."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (1, 2, 3), (1, 0, 0, 0))
        c = CoincidentConstraint(b1, (0, 0, 0), b2, (0, 0, 0))
        env = pt.get_env()
        vals = [r.eval(env) for r in c.residuals()]
        assert abs(vals[0] - (-1.0)) < 1e-10
        assert abs(vals[1] - (-2.0)) < 1e-10
        assert abs(vals[2] - (-3.0)) < 1e-10

    def test_with_markers(self):
        """Marker offsets shift the coincidence point."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (10, 0, 0), (1, 0, 0, 0))
        # marker on b1 at local (5,0,0), marker on b2 at local (-5,0,0)
        c = CoincidentConstraint(b1, (5, 0, 0), b2, (-5, 0, 0))
        env = pt.get_env()
        # b1 marker world = (5,0,0), b2 marker world = (10-5, 0, 0) = (5,0,0)
        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 = CoincidentConstraint(b1, (0, 0, 0), b2, (0, 0, 0))
        assert len(c.residuals()) == 3


class TestDistancePointPoint:
    def test_satisfied(self):
        """Bodies at distance 5 apart with d=5 → residual zero."""
        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 = DistancePointPointConstraint(b1, (0, 0, 0), b2, (0, 0, 0), 5.0)
        env = pt.get_env()
        # 3^2 + 4^2 - 5^2 = 9+16-25 = 0
        assert abs(c.residuals()[0].eval(env)) < 1e-10

    def test_unsatisfied(self):
        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 = DistancePointPointConstraint(b1, (0, 0, 0), b2, (0, 0, 0), 3.0)
        env = pt.get_env()
        # 9+16-9 = 16
        assert abs(c.residuals()[0].eval(env) - 16.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 = DistancePointPointConstraint(b1, (0, 0, 0), b2, (0, 0, 0), 1.0)
        assert len(c.residuals()) == 1


class TestFixed:
    def test_satisfied(self):
        """Identical pose + markers → all residuals zero."""
        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 = FixedConstraint(
            b1,
            (0, 0, 0),
            (1, 0, 0, 0),
            b2,
            (0, 0, 0),
            (1, 0, 0, 0),
        )
        env = pt.get_env()
        for r in c.residuals():
            assert abs(r.eval(env)) < 1e-10

    def test_residual_count(self):
        """Fixed constraint: 3 position + 3 orientation = 6 residuals."""
        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 = FixedConstraint(
            b1,
            (0, 0, 0),
            (1, 0, 0, 0),
            b2,
            (0, 0, 0),
            (1, 0, 0, 0),
        )
        assert len(c.residuals()) == 6

    def test_unsatisfied_position(self):
        """Offset position → non-zero position residuals."""
        pt = ParamTable()
        b1 = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        b2 = RigidBody("b", pt, (1, 0, 0), (1, 0, 0, 0))
        c = FixedConstraint(
            b1,
            (0, 0, 0),
            (1, 0, 0, 0),
            b2,
            (0, 0, 0),
            (1, 0, 0, 0),
        )
        env = pt.get_env()
        vals = [r.eval(env) for r in c.residuals()]
        # Position residual [0] should be -1 (0 - 1)
        assert abs(vals[0] - (-1.0)) < 1e-10