"""Regression tests for interactive drag.

These tests exercise the drag protocol at the solver-internals level,
verifying that constraints remain enforced across drag steps when the
pre-pass has been applied to cached residuals.

Bug scenario: single_equation_pass runs during pre_drag, analytically
solving variables from upstream constraints and baking their values as
constants into downstream residual expressions.  When a drag step
changes those variables, the cached residuals use stale constants and
downstream constraints (e.g. Planar distance=0) stop being enforced.

Fix: skip single_equation_pass in the drag path.  Only substitution_pass
(which replaces genuinely grounded parameters) is safe to cache.
"""

import math

import pytest
from kindred_solver.constraints import (
    CoincidentConstraint,
    PlanarConstraint,
    RevoluteConstraint,
)
from kindred_solver.entities import RigidBody
from kindred_solver.newton import newton_solve
from kindred_solver.params import ParamTable
from kindred_solver.prepass import single_equation_pass, substitution_pass

ID_QUAT = (1, 0, 0, 0)


def _build_residuals(bodies, constraint_objs):
    """Build raw residual list + quat groups (no prepass)."""
    all_residuals = []
    for c in constraint_objs:
        all_residuals.extend(c.residuals())

    quat_groups = []
    for body in bodies:
        if not body.grounded:
            all_residuals.append(body.quat_norm_residual())
            quat_groups.append(body.quat_param_names())

    return all_residuals, quat_groups


def _eval_raw_residuals(bodies, constraint_objs, params):
    """Evaluate original constraint residuals at current param values.

    Returns the max absolute residual value — the ground truth for
    whether constraints are satisfied regardless of prepass state.
    """
    raw, _ = _build_residuals(bodies, constraint_objs)
    env = params.get_env()
    return max(abs(r.eval(env)) for r in raw)


class TestPrepassDragRegression:
    """single_equation_pass bakes stale values that break drag.

    Setup: ground --Revolute--> arm --Planar(d=0)--> plate

    The Revolute pins arm's origin to ground (fixes arm/tx, arm/ty,
    arm/tz via single_equation_pass).  The Planar keeps plate coplanar
    with arm.  After prepass, the Planar residual has arm's position
    baked as Const(0.0).

    During drag: arm/tz is set to 5.0.  Because arm/tz is marked fixed
    by prepass, Newton can't correct it, AND the Planar residual still
    uses Const(0.0) instead of the live value 5.0.  The Revolute
    constraint (arm at origin) is silently violated.
    """

    def _setup(self):
        pt = ParamTable()
        ground = RigidBody("g", pt, (0, 0, 0), ID_QUAT, grounded=True)
        arm = RigidBody("arm", pt, (10, 0, 0), ID_QUAT)
        plate = RigidBody("plate", pt, (10, 5, 0), ID_QUAT)

        constraints = [
            RevoluteConstraint(ground, (0, 0, 0), ID_QUAT, arm, (0, 0, 0), ID_QUAT),
            PlanarConstraint(arm, (0, 0, 0), ID_QUAT, plate, (0, 0, 0), ID_QUAT, offset=0.0),
        ]
        bodies = [ground, arm, plate]
        return pt, bodies, constraints

    def test_bug_stale_constants_after_single_equation_pass(self):
        """Document the bug: prepass bakes arm/tz=0, drag breaks constraints."""
        pt, bodies, constraints = self._setup()
        raw_residuals, quat_groups = _build_residuals(bodies, constraints)

        # Simulate OLD pre_drag: substitution + single_equation_pass
        residuals = substitution_pass(raw_residuals, pt)
        residuals = single_equation_pass(residuals, pt)

        ok = newton_solve(residuals, pt, quat_groups=quat_groups, max_iter=100, tol=1e-10)
        assert ok

        # Verify prepass fixed arm's position params
        assert pt.is_fixed("arm/tx")
        assert pt.is_fixed("arm/ty")
        assert pt.is_fixed("arm/tz")

        # Simulate drag: move arm up (set_value, as drag_step does)
        pt.set_value("arm/tz", 5.0)
        pt.set_value("plate/tz", 5.0)  # initial guess near drag

        ok = newton_solve(residuals, pt, quat_groups=quat_groups, max_iter=100, tol=1e-10)
        # Solver "converges" on the stale cached residuals
        assert ok

        # But the TRUE constraints are violated: arm should be at z=0
        # (Revolute pins it to ground) yet it's at z=5
        max_err = _eval_raw_residuals(bodies, constraints, pt)
        assert max_err > 1.0, (
            f"Expected large raw residual violation, got {max_err:.6e}. "
            "The bug should cause the Revolute z-residual to be ~5.0"
        )

    def test_fix_no_single_equation_pass_for_drag(self):
        """With the fix: skip single_equation_pass, constraints hold."""
        pt, bodies, constraints = self._setup()
        raw_residuals, quat_groups = _build_residuals(bodies, constraints)

        # Simulate FIXED pre_drag: substitution only
        residuals = substitution_pass(raw_residuals, pt)

        ok = newton_solve(residuals, pt, quat_groups=quat_groups, max_iter=100, tol=1e-10)
        assert ok

        # arm/tz should NOT be fixed
        assert not pt.is_fixed("arm/tz")

        # Simulate drag: move arm up
        pt.set_value("arm/tz", 5.0)
        pt.set_value("plate/tz", 5.0)

        ok = newton_solve(residuals, pt, quat_groups=quat_groups, max_iter=100, tol=1e-10)
        assert ok

        # Newton pulls arm back to z=0 (Revolute enforced) and plate follows
        max_err = _eval_raw_residuals(bodies, constraints, pt)
        assert max_err < 1e-8, f"Raw residual violation {max_err:.6e} — constraints not satisfied"


class TestCoincidentPlanarDragRegression:
    """Coincident upstream + Planar downstream — same bug class.

    ground --Coincident--> bracket --Planar(d=0)--> plate

    Coincident fixes bracket/tx,ty,tz.  After prepass, the Planar
    residual has bracket's position baked.  Drag moves bracket;
    the Planar uses stale constants.
    """

    def _setup(self):
        pt = ParamTable()
        ground = RigidBody("g", pt, (0, 0, 0), ID_QUAT, grounded=True)
        bracket = RigidBody("bracket", pt, (0, 0, 0), ID_QUAT)
        plate = RigidBody("plate", pt, (10, 5, 0), ID_QUAT)

        constraints = [
            CoincidentConstraint(ground, (0, 0, 0), bracket, (0, 0, 0)),
            PlanarConstraint(bracket, (0, 0, 0), ID_QUAT, plate, (0, 0, 0), ID_QUAT, offset=0.0),
        ]
        bodies = [ground, bracket, plate]
        return pt, bodies, constraints

    def test_bug_coincident_planar(self):
        """Prepass fixes bracket/tz, Planar uses stale constant during drag."""
        pt, bodies, constraints = self._setup()
        raw, qg = _build_residuals(bodies, constraints)

        residuals = substitution_pass(raw, pt)
        residuals = single_equation_pass(residuals, pt)

        ok = newton_solve(residuals, pt, quat_groups=qg, max_iter=100, tol=1e-10)
        assert ok
        assert pt.is_fixed("bracket/tz")

        # Drag bracket up
        pt.set_value("bracket/tz", 5.0)
        pt.set_value("plate/tz", 5.0)

        ok = newton_solve(residuals, pt, quat_groups=qg, max_iter=100, tol=1e-10)
        assert ok

        # True constraints violated
        max_err = _eval_raw_residuals(bodies, constraints, pt)
        assert max_err > 1.0, f"Expected raw violation from stale prepass, got {max_err:.6e}"

    def test_fix_coincident_planar(self):
        """With the fix: constraints satisfied after drag."""
        pt, bodies, constraints = self._setup()
        raw, qg = _build_residuals(bodies, constraints)

        residuals = substitution_pass(raw, pt)
        # No single_equation_pass

        ok = newton_solve(residuals, pt, quat_groups=qg, max_iter=100, tol=1e-10)
        assert ok
        assert not pt.is_fixed("bracket/tz")

        # Drag bracket up
        pt.set_value("bracket/tz", 5.0)
        pt.set_value("plate/tz", 5.0)

        ok = newton_solve(residuals, pt, quat_groups=qg, max_iter=100, tol=1e-10)
        assert ok

        max_err = _eval_raw_residuals(bodies, constraints, pt)
        assert max_err < 1e-8, f"Raw residual violation {max_err:.6e} — constraints not satisfied"


class TestDragDoesNotBreakStaticSolve:
    """Verify that the static solve path (with single_equation_pass) still works.

    The fix only affects pre_drag — the static solve() path continues to
    use single_equation_pass for faster convergence.
    """

    def test_static_solve_still_uses_prepass(self):
        """Static solve with single_equation_pass converges correctly."""
        pt = ParamTable()
        ground = RigidBody("g", pt, (0, 0, 0), ID_QUAT, grounded=True)
        arm = RigidBody("arm", pt, (10, 0, 0), ID_QUAT)
        plate = RigidBody("plate", pt, (10, 5, 8), ID_QUAT)

        constraints = [
            RevoluteConstraint(ground, (0, 0, 0), ID_QUAT, arm, (0, 0, 0), ID_QUAT),
            PlanarConstraint(arm, (0, 0, 0), ID_QUAT, plate, (0, 0, 0), ID_QUAT, offset=0.0),
        ]
        bodies = [ground, arm, plate]
        raw, qg = _build_residuals(bodies, constraints)

        # Full prepass (static solve path)
        residuals = substitution_pass(raw, pt)
        residuals = single_equation_pass(residuals, pt)

        ok = newton_solve(residuals, pt, quat_groups=qg, max_iter=100, tol=1e-10)
        assert ok

        # All raw constraints satisfied
        max_err = _eval_raw_residuals(bodies, constraints, pt)
        assert max_err < 1e-8

        # arm at origin (Revolute), plate coplanar (z=0)
        env = pt.get_env()
        assert abs(env["arm/tx"]) < 1e-8
        assert abs(env["arm/ty"]) < 1e-8
        assert abs(env["arm/tz"]) < 1e-8
        assert abs(env["plate/tz"]) < 1e-8