"""The eigensolver: all solvers agree, mass normalization and residuals, degenerate
pairs of a disc, and the free square and disc frequencies against Gorman's and
Leissa's values."""

import numpy as np
import pytest

from cymsim import Plate, compute_modes, make_grid
from cymsim.linear.assemble import assemble

from conftest import STEEL


def test_solvers_agree_and_shapes_are_mass_normalized():
    pl = Plate.uniform(make_grid(0.12, 0.12, 0.012, lambda x, y, r, t: r <= 0.06), STEEL, 1e-3)
    sys = assemble(pl)
    a = compute_modes(pl, 20000, solver="dense", damping=False, system=sys)
    b = compute_modes(pl, 20000, solver="banded", damping=False, system=sys)
    c = compute_modes(pl, 20000, solver="torch", device="cpu", damping=False, system=sys)
    assert a.count == b.count == c.count > 10
    assert np.allclose(a.freqs, b.freqs, rtol=1e-8) and np.allclose(a.freqs, c.freqs, rtol=1e-8)
    for m in (0, 1, a.count // 2, a.count - 1):
        phi = sys.from_grid(np.zeros(pl.grid.shape), np.zeros(pl.grid.shape), a.w[m])
        assert abs(phi @ (sys.mass * phi) - 1) < 1e-9
        lam = (2 * np.pi * a.freqs[m]) ** 2
        r = sys.K @ phi - lam * sys.mass * phi
        assert np.linalg.norm(r) < 1e-8 * np.linalg.norm(lam * sys.mass * phi)
    # The disc has degenerate pairs, which the staircase splits by about a percent.
    assert abs(a.freqs[1] / a.freqs[0] - 1) < 0.02
    # No rigid modes survive.
    assert a.freqs[0] > 1.0


def test_select_and_below():
    pl = Plate.uniform(make_grid(0.1, 0.08, 0.01), STEEL, 1e-3)
    ms = compute_modes(pl, 20000, damping=False)
    sub = ms.below(5000)
    assert sub.count == (ms.freqs < 5000).sum() and sub.w.shape[0] == sub.count
    two = ms.select([0, 2])
    assert np.array_equal(two.freqs, ms.freqs[[0, 2]])


def test_free_square_frequencies_converge_to_gorman():
    # lambda = omega a^2 sqrt(rho h / D) for the completely free square (nu = 0.3).
    gorman = np.array([13.468, 19.596, 24.271, 34.801, 34.801, 61.093, 61.093])
    a, h = 0.3, 1e-3
    scale = a * a * np.sqrt(STEEL.rho * h / STEEL.bending_stiffness(h))
    errors = []
    for cells in (10, 20):
        ms = compute_modes(Plate.uniform(make_grid(a, a, a / cells), STEEL, h), 80 / scale / (2 * np.pi), damping=False)
        errors.append(np.abs(2 * np.pi * ms.freqs[:7] * scale / gorman - 1).max())
    assert errors[0] < 0.1 and errors[1] < 0.025 and errors[1] < errors[0]


def test_free_disc_frequencies_are_close_to_leissa():
    # lambda^2 = omega a^2 sqrt(rho h / D), completely free disc, nu = 0.3.
    leissa = np.array([5.253, 5.253, 9.084, 12.23, 12.23, 20.52, 20.52, 21.6, 21.6])
    a, h = 0.15, 1e-3
    scale = a * a * np.sqrt(STEEL.rho * h / STEEL.bending_stiffness(h))
    g = make_grid(2 * a, 2 * a, 2 * a / 40, lambda x, y, r, t: r <= a)
    ms = compute_modes(Plate.uniform(g, STEEL, h), 25 / scale / (2 * np.pi), damping=False)
    assert np.abs(2 * np.pi * ms.freqs[:9] * scale / leissa - 1).max() < 0.03


def test_save_and_load(tmp_path):
    pl = Plate.uniform(make_grid(0.1, 0.08, 0.01), STEEL, 1e-3, name="small")
    ms = compute_modes(pl, 8000)
    ms.save(tmp_path / "m.npz")
    back = ms.__class__.load(tmp_path / "m.npz")
    assert np.array_equal(back.freqs, ms.freqs) and np.array_equal(back.w, ms.w)
    assert back.plate.name == "small" and back.plate.material == pl.material
    assert np.array_equal(back.sigma_radiation, ms.sigma_radiation)


@pytest.mark.slow
def test_crash_16_mode_count():
    from cymsim import presets

    pl = presets.crash_16().evaluate(0.008)
    ms = compute_modes(pl, 16000)
    # About 1900 bending modes below 16 kHz on this grid (Weyl's estimate gives the same order).
    assert 1700 < ms.count < 2100
    assert ms.freqs[0] > 15 and ms.freqs[0] < 40