from pathlib import Path

import numpy as np

from femsurrogate.fea.assembly import assemble_global_stiffness, constrained_dofs
from femsurrogate.fea.io import read_beam_example
from femsurrogate.fea.solver import solve_linear_static

ROOT = Path(__file__).resolve().parents[1]


def test_assemble_global_stiffness_for_cantilever_fixture_has_expected_shape():
    model = read_beam_example(ROOT / "BeamExamples" / "CantileverBeam.txt")
    stiffness = assemble_global_stiffness(model)

    assert stiffness.shape == (18, 18)


def test_fixed_node_dofs_are_constrained():
    model = read_beam_example(ROOT / "BeamExamples" / "CantileverBeam.txt")

    assert constrained_dofs(model) == (0, 1, 2)


def test_solve_linear_static_returns_finite_displacements_for_all_nodes():
    model = read_beam_example(ROOT / "BeamExamples" / "CantileverBeam.txt")
    displacements = solve_linear_static(model)

    assert set(displacements) == set(model.nodes)
    values = np.array([[d.ux, d.uy, d.rz] for d in displacements.values()])
    assert np.isfinite(values).all()