import numpy as np
import pytest

from femsurrogate.fea.element import local_frame_stiffness, transformation_matrix


def test_local_frame_stiffness_has_expected_shape_and_symmetry():
    stiffness = local_frame_stiffness(E=200.0, A=3.0, Izz=5.0, L=7.0)

    assert stiffness.shape == (6, 6)
    np.testing.assert_allclose(stiffness, stiffness.T)


def test_local_frame_stiffness_contains_clockwise_rotation_terms():
    E = 210.0
    A = 2.0
    Izz = 4.0
    L = 3.0

    stiffness = local_frame_stiffness(E=E, A=A, Izz=Izz, L=L)

    assert stiffness[0, 0] == pytest.approx(E * A / L)
    assert stiffness[0, 3] == pytest.approx(-E * A / L)
    assert stiffness[1, 1] == pytest.approx(12 * E * Izz / L**3)
    assert stiffness[1, 2] == pytest.approx(-6 * E * Izz / L**2)
    assert stiffness[2, 2] == pytest.approx(4 * E * Izz / L)
    assert stiffness[2, 5] == pytest.approx(2 * E * Izz / L)
    assert stiffness[4, 4] == pytest.approx(12 * E * Izz / L**3)
    assert stiffness[4, 5] == pytest.approx(6 * E * Izz / L**2)
    assert stiffness[5, 5] == pytest.approx(4 * E * Izz / L)


def test_transformation_matrix_is_identity_for_horizontal_element():
    matrix = transformation_matrix(0.0, 0.0, 2.0, 0.0)

    np.testing.assert_allclose(matrix, np.eye(6))


def test_transformation_matrix_is_orthogonal_for_inclined_element():
    matrix = transformation_matrix(0.0, 0.0, 1.0, 1.0)

    np.testing.assert_allclose(matrix.T @ matrix, np.eye(6), atol=1e-12)