"""Tests for the finite-element (linear elasticity) deformation model.

Small domains keep each FEM solve fast. The checks assert the physical
properties a credible elastostatic solve must have — radial outward
displacement, monotonic decay to the fixed far field, symmetry for an
axis-aligned probe, and drop-in compatibility with the risk pipeline.
"""
import numpy as np
import pytest

from sulcora.volume import Volume
from sulcora.planning import Trajectory
from sulcora.fem import FEMProbeDeformation, _hex_stiffness


def _fem(**kw):
    traj = Trajectory(entry=(0, 0, -10), target=(0, 0, 10))   # probe along +z
    params = dict(probe_radius_mm=2.0, influence_radius_mm=8.0, element_mm=2.5)
    params.update(kw)
    return FEMProbeDeformation(traj, **params)


def test_element_stiffness_is_symmetric_psd():
    Ke = _hex_stiffness(2.5, 0.45)
    assert Ke.shape == (24, 24)
    assert np.allclose(Ke, Ke.T, atol=1e-9)                 # symmetric
    # 6 rigid-body modes -> exactly 6 (near-)zero eigenvalues, rest positive
    w = np.linalg.eigvalsh(Ke)
    assert np.sum(w < 1e-6 * w.max()) == 6
    assert np.all(w > -1e-6 * w.max())                      # positive semidefinite


def test_radial_outward_and_decay():
    fem = _fem()
    mags = [np.linalg.norm(fem.field(np.array([[r, 0.0, 0.0]]))[0])
            for r in (2.5, 5.0, 7.5)]
    assert mags[0] > mags[1] > mags[2] > 0.0                # monotonic decay
    v = fem.field(np.array([[5.0, 0.0, 0.0]]))[0]
    assert v[0] > 0.0                                       # outward (+x)
    assert abs(v[2]) < 0.05 * np.linalg.norm(v) + 1e-6      # ~no axial component


def test_axis_aligned_probe_is_symmetric():
    fem = _fem()
    m = lambda p: np.linalg.norm(fem.field(np.array([p], float))[0])
    a, b, c, d = m((5, 0, 0)), m((-5, 0, 0)), m((0, 5, 0)), m((0, -5, 0))
    assert np.allclose([a, a, a], [b, c, d], rtol=1e-6, atol=1e-9)


def test_max_shift_bounded_and_finite():
    fem = _fem(probe_radius_mm=2.0)
    fem._ensure_solved()
    assert np.all(np.isfinite(fem._U))
    assert fem.max_shift_mm == 2.0
    # nowhere should the solved displacement exceed the prescribed wall motion
    assert float(np.linalg.norm(fem._U, axis=-1).max()) <= 2.0 + 1e-6


def test_zero_outside_domain():
    fem = _fem()
    far = fem.field(np.array([[500.0, 500.0, 500.0]]))[0]
    assert np.allclose(far, 0.0)


def test_poisson_changes_field():
    near = lambda nu: np.linalg.norm(_fem(poisson=nu).field(np.array([[5.0, 0, 0]]))[0])
    # different Poisson ratio -> different elastic response
    assert abs(near(0.30) - near(0.49)) > 1e-3


def test_drop_in_to_deformation_aware_risk():
    from sulcora import risk as R
    aff = np.eye(4)
    data = np.zeros((31, 31, 31), np.int16)
    data[5, 15, 15] = 2                                     # "Motor" at world (5,15,15)
    av = Volume(data, aff)
    labels = ["Background", "Other", "Motor"]
    traj = Trajectory(entry=(9, 0, 15), target=(9, 30, 15))  # 4mm from structure
    fem = FEMProbeDeformation(traj, probe_radius_mm=2.0, influence_radius_mm=12.0, element_mm=3.0)
    rep = R.deformation_aware_risk(traj, av, labels, names=["Motor"], deformation=fem, n=201)
    assert abs(rep["static_clearance_mm"] - 4.0) < 0.2
    assert "deformed_clearance_mm" in rep                   # FEM is a valid deformation source
    assert rep["deformed_clearance_mm"] != rep["static_clearance_mm"]


def test_summary_keys():
    s = _fem().summary()
    assert s["model"] == "FEM linear elasticity"
    assert "max_shift_mm" in s and s["poisson"] == 0.45
