import jax
import jax.numpy as jnp
import numpy as np

from cymsim.elements import dkt_stiffness, element_energy_parts, membrane_strains


def _triangle():
    return jnp.array([0.0, 0.03, 0.01]), jnp.array([0.0, 0.005, 0.025])


def test_dkt_rigid_modes_and_constant_curvature():
    xe, ye = _triangle()
    D, nu = 10.0, 0.3
    Ke = np.asarray(dkt_stiffness(xe, ye, D, nu))
    assert np.allclose(Ke, Ke.T)
    x, y = np.asarray(xe), np.asarray(ye)
    # rigid modes: translation, tilt about x (w = y, thx = 1), tilt about y (w = -x, thy = 1)
    for U in (np.tile([1, 0, 0], 3), np.stack([y, np.ones(3), np.zeros(3)], 1).ravel(),
              np.stack([-x, np.zeros(3), np.ones(3)], 1).ravel()):
        assert np.abs(Ke @ U).max() < 1e-12 * np.abs(Ke).max()
    # constant curvature w = x^2 / 2: energy D A / 2
    A = 0.5 * ((x[1] - x[0]) * (y[2] - y[0]) - (x[2] - x[0]) * (y[1] - y[0]))
    U = np.stack([0.5 * x**2, np.zeros(3), -x], 1).ravel()
    assert np.isclose(0.5 * U @ Ke @ U, 0.5 * D * A, rtol=1e-9)
    # twist w = x y: energy D (1 - nu) A
    U = np.stack([x * y, x, -y], 1).ravel()
    assert np.isclose(0.5 * U @ Ke @ U, D * (1 - nu) * A, rtol=1e-9)
    eig = np.linalg.eigvalsh(Ke)
    assert (eig[:3] < 1e-10 * eig[-1]).all() and eig[3] > 1e-6 * eig[-1]


def test_membrane_strains_linear_and_nonlinear():
    Gx, Gy = jnp.array([1.0, -1.0, 0.0]), jnp.array([0.0, 1.0, -1.0])
    # nodal values arranged so that u_x = 2, v_y = 3, w_x = 0.1, w_y = 0.2 for these operators
    qe = jnp.zeros((3, 5)).at[:, 0].set(jnp.array([2.0, 0.0, 0.0])).at[:, 1].set(jnp.array([0.0, 3.0, 0.0])) \
        .at[:, 2].set(jnp.array([0.1, 0.0, -0.2]))
    exx, eyy, gxy = membrane_strains(qe, Gx, Gy, 0.5, -0.5, 1.0)
    assert np.isclose(exx, 2 + 0.5 * 0.01 + 0.5 * 0.1)
    assert np.isclose(eyy, 3 + 0.5 * 0.04 - 0.5 * 0.2)
    assert np.isclose(gxy, -3 + 0.1 * 0.2 + 0.5 * 0.2 - 0.5 * 0.1)   # u_y + v_x = 0 - 3
    exx0, eyy0, _ = membrane_strains(qe, Gx, Gy, 0.5, -0.5, 0.0)
    assert np.isclose(exx0, 2 + 0.5 * 0.1) and np.isclose(eyy0, 3 - 0.5 * 0.2)


def test_energy_hessian_is_symmetric_positive_semidefinite():
    xe, ye = _triangle()
    Ke = dkt_stiffness(xe, ye, 10.0, 0.3)
    x, y = np.asarray(xe), np.asarray(ye)
    A = 0.5 * ((x[1] - x[0]) * (y[2] - y[0]) - (x[2] - x[0]) * (y[1] - y[0]))
    Gx = jnp.array([y[1] - y[2], y[2] - y[0], y[0] - y[1]]) / (2 * A)
    Gy = jnp.array([x[2] - x[1], x[0] - x[2], x[1] - x[0]]) / (2 * A)
    Cm = jnp.array([1e6, 3e5, 3.5e5])
    f = lambda q: sum(element_energy_parts(q, Gx, Gy, A, 0.1, -0.05, Cm, Ke, 1.0))
    H = np.asarray(jax.hessian(f)(jnp.zeros((3, 5)))).reshape(15, 15)
    assert np.allclose(H, H.T)
    eig = np.linalg.eigvalsh(H)
    assert eig[0] > -1e-9 * eig[-1]
    assert (eig < 1e-9 * eig[-1]).sum() == 6      # 3 translations + 3 rotations