import numpy as np
import pytest

from cymsim.config import Material, Numerics, Profile, Imperfections
from cymsim.geometry import CymbalProfile
from cymsim.mesh import disk_mesh
from cymsim.shell import ShellModel


@pytest.fixture(scope="module")
def shell():
    prof = CymbalProfile(Profile(), Imperfections(n_dents=30))
    mesh = disk_mesh(Profile.radius, 0.025)
    return ShellModel(mesh, prof, Material(), Numerics(dtype="float64"))


def test_mass_matches_integrated_density(shell):
    rho = shell.material.rho
    expected = rho * (shell.h_e * shell.mesh.area).sum()
    assert np.isclose(shell.mass, expected)
    assert (shell.M > 0).all()


def test_element_checks(shell):
    c = shell.element_checks()
    assert c.rigid_tilt_residual < 1e-12
    assert c.constant_curvature_error < 1e-8
    assert c.zero_energy_modes_min == 6 and c.zero_energy_modes_max == 6
    assert c.rigid_mode_ratio < 1e-10


def test_rigid_modes_are_M_orthonormal_and_force_free(shell):
    import jax
    import jax.numpy as jnp
    G = np.einsum("knd,nd,lnd->kl", shell.Rm, shell.M, shell.Rm)
    assert np.allclose(G, np.eye(6), atol=1e-10)
    # the tangent stiffness at rest annihilates every rigid mode
    q0 = jnp.zeros(shell.M.shape)
    K = lambda d: np.asarray(jax.jvp(shell.internal_force, (q0,), (jnp.asarray(d),))[1])
    rng = np.random.default_rng(0)
    for k in range(6):
        rnd = rng.standard_normal(shell.M.shape)
        rnd *= np.linalg.norm(shell.Rm[k]) / np.linalg.norm(rnd)
        assert np.abs(K(shell.Rm[k])).max() < 1e-12 * np.abs(K(rnd)).max()
    # a finite tilt is force-free only to first order: the residual scales with amplitude squared
    r1 = np.abs(np.asarray(shell.internal_force(shell.Rm_j[4] * 1e-3))).max()
    r2 = np.abs(np.asarray(shell.internal_force(shell.Rm_j[4] * 2e-3))).max()
    assert 3.5 < r2 / r1 < 4.5


def test_projection_removes_rigid_motion(shell):
    rng = np.random.default_rng(1)
    v = rng.standard_normal(shell.M.shape)
    p = np.asarray(shell.project_rigid(v))
    assert np.allclose(np.einsum("knd,nd->k", shell.Rm, shell.M * p), 0, atol=1e-9)
    assert np.allclose(np.asarray(shell.project_rigid(p)), p)


def test_internal_force_is_minus_energy_gradient(shell):
    rng = np.random.default_rng(2)
    q = 1e-4 * rng.standard_normal(shell.M.shape)
    d = rng.standard_normal(shell.M.shape)
    eps = 1e-6
    dE = (float(shell.total_energy(q + eps * d)) - float(shell.total_energy(q - eps * d))) / (2 * eps)
    assert np.isclose(-(np.asarray(shell.internal_force(q)) * d).sum(), dE, rtol=1e-5)


def test_damping_is_dissipative(shell):
    rng = np.random.default_rng(3)
    vw = rng.standard_normal(shell.mesh.n_nodes)
    fd = np.asarray(shell.damping_force_w(vw))
    assert (vw * fd).sum() > 0        # the force -fd does negative work on vw
    assert np.isclose(fd.sum(), 0.0, atol=1e-9 * np.abs(fd).sum())   # a uniform velocity is undamped


def test_stability_bound_is_finite_and_sane(shell):
    assert 1e4 < shell.omega_max < 1e8
    assert shell.omega_global < shell.omega_max          # the element bound is an upper bound
    assert shell.dt_stable == 2.0 / shell.omega_global   # the default numerics.stability_bound is "global"


def test_global_bound_matches_the_dense_eigenproblem(shell):
    """The power iteration finds the largest eigenvalue of M^-1 K, K assembled from the tangents at rest."""
    N, L = shell.mesh.n_nodes, shell.M.size
    K = np.zeros((L, L))
    dof = (5 * shell.mesh.tri[:, :, None] + np.arange(5)).reshape(-1, 15)
    for e in range(len(dof)):
        K[np.ix_(dof[e], dof[e])] += shell.K_rest[e]
    isq = 1 / np.sqrt(shell.M.ravel())
    lam = np.linalg.eigvalsh(isq[:, None] * K * isq[None, :])[-1]
    assert abs(shell.omega_global / np.sqrt(lam) - 1) < 1e-6


def test_element_bound_on_request(shell):
    sh = ShellModel(shell.mesh, CymbalProfile(Profile(), Imperfections(n_dents=30)), Material(),
                    Numerics(dtype="float64", stability_bound="element"))
    assert sh.dt_stable == 2.0 / sh.omega_max