"""The discretization: rigid modes of the stiffness, the assembled K against the
linearized kernel, forces as the gradient of the energy, dissipativity of the
losses, and thickness scaling."""

import numpy as np
import pytest
import torch

from cymsim import Plate, make_grid
from cymsim.full.kernel import accelerations, damping, static_data
from cymsim.full.reference import max_angular_frequency, potential_energy
from cymsim.linear.assemble import assemble

from conftest import BRONZE, STEEL

T = lambda a: torch.as_tensor(np.asarray(a, dtype=np.float64))  # noqa: E731


def bowed(g, sag):
    Lx = (g.nx - 1) * g.hx
    return sag * (1 - (2 * g.X / Lx) ** 2) + 0.3 * sag * np.sin(3 * g.Y / Lx)


def random_fields(pl, amp, seed):
    r = np.random.default_rng(seed)
    return [np.where(pl.grid.mask, a * (r.random(pl.grid.shape) - 0.5), 0.0) for a in (0.02 * amp, 0.02 * amp, amp)]


@pytest.mark.parametrize("case", ["uniform", "tapered", "disc-hole"])
def test_stiffness_has_exactly_the_rigid_modes_in_its_null_space(case):
    if case == "disc-hole":
        g = make_grid(0.16, 0.16, 0.01, lambda x, y, r, t: (r <= 0.08) & (r >= 0.015))
        pl, fields = Plate.uniform(g, STEEL, 1e-3), 1
    else:
        g = make_grid(0.12, 0.09, 0.015)
        pl, fields = Plate.uniform(g, STEEL, 1e-3), 3
        if case == "tapered":
            pl = pl.with_thickness(1e-3 * (1 + 0.5 * (g.X + 0.06) / 0.12))
    sys = assemble(pl, fields)
    zero = np.zeros(g.shape)
    one = np.ones(g.shape)
    null = [(zero, zero, one), (zero, zero, g.X), (zero, zero, g.Y)]
    if fields == 3:
        null += [(one, zero, zero), (zero, one, zero), (-g.Y, g.X, zero)]
    scale = abs(sys.K[0, 0])
    for u, v, w in null:
        x = sys.from_grid(u, v, w)
        assert np.linalg.norm(sys.K @ x) < 1e-9 * scale * np.linalg.norm(x)
    x = sys.from_grid(zero, zero, np.sin(0.9 * np.arange(g.nx))[None, :] * np.cos(1.1 * np.arange(g.ny))[:, None])
    assert np.abs(sys.K @ x).max() > 1e-3
    assert abs(sys.K - sys.K.T).max() < 1e-9 * scale


@pytest.mark.parametrize("fields", [1, 3])
def test_assembled_K_equals_the_linearized_kernel_on_a_flat_plate(fields):
    pl = Plate.uniform(make_grid(0.2, 0.15, 0.02), BRONZE, 1e-3)
    assert pl.is_flat
    assert relative_error_against_kernel(pl, fields) < 1e-6


def test_assembled_K_equals_the_linearized_kernel_on_a_bowed_tapered_plate_with_a_hole():
    g = make_grid(0.2, 0.2, 0.0125, lambda x, y, r, t: (r <= 0.1) & (r >= 0.01))
    pl = Plate.uniform(g, BRONZE, 1e-3, profile=3e-3 * np.sin(20 * g.X) * np.cos(15 * g.Y))
    pl = pl.with_thickness(1e-3 * (1 + 0.4 * np.abs(g.X) / 0.1))
    assert not pl.is_flat
    assert assemble(pl).fields == 3
    assert relative_error_against_kernel(pl, 3) < 1e-6


def relative_error_against_kernel(pl, fields, seed=3):
    sys = assemble(pl, fields)
    x = np.random.default_rng(seed).random(sys.n) - 0.5
    y = sys.K @ x
    eps = 1e-9
    u, v, w = sys.to_grid(x * eps)
    S = static_data(pl, None, 1e-6, 0, 0, 0, "cpu", torch.float64)
    au, av, aw = (a.numpy() for a in accelerations(S, T(u), T(v), T(w)))
    m = pl.node_mass
    yref = sys.from_grid(-au * m / eps, -av * m / eps, -aw * m / eps)
    return np.linalg.norm(y - yref) / np.linalg.norm(yref)


@pytest.mark.parametrize("case", ["flat", "bowed", "bowed-tapered"])
def test_forces_are_the_gradient_of_the_energy(case):
    g = make_grid(0.1, 0.08, 0.01)
    pl = Plate.uniform(g, STEEL, 1e-3, profile=None if case == "flat" else bowed(g, 3e-3))
    if case == "bowed-tapered":
        pl = pl.with_thickness(1e-3 * (1 + 0.4 * (g.X + 0.05) / 0.1))
    f = random_fields(pl, 2.5e-3, 7)
    d = random_fields(pl, 1.0, 11)
    S = static_data(pl, None, 1e-6, 0, 0, 0, "cpu", torch.float64)
    acc = [a.numpy() for a in accelerations(S, *[T(a) for a in f])]
    h = 1e-7
    dU = (potential_energy(pl, *[a + h * b for a, b in zip(f, d)]) - potential_energy(pl, *[a - h * b for a, b in zip(f, d)])) / (2 * h)
    fd = -sum((pl.node_mass * a * b).sum() for a, b in zip(acc, d))
    assert abs(dU - fd) < 1e-6 * abs(dU)


def test_the_loss_terms_only_remove_energy():
    pl = Plate.uniform(make_grid(0.1, 0.08, 0.01), STEEL, 1e-3)
    S = static_data(pl, None, 1e-6, 1.0, 5e-4, 20.0, "cpu", torch.float64)
    for seed in (1, 2, 3):
        vel = random_fields(pl, 1.0, seed)
        acc = [a.numpy() for a in damping(S, *[T(a) for a in vel])]
        power = sum((pl.node_mass * a * b).sum() for a, b in zip(acc, vel))
        assert power < 0


def test_doubling_the_thickness_doubles_the_bending_frequencies_only():
    g = make_grid(0.1, 0.1, 0.01)
    w1, w2 = max_angular_frequency(g, STEEL, 1e-3), max_angular_frequency(g, STEEL, 2e-3)
    # The largest frequency is set by the in-plane waves, independent of h.
    assert abs(w1 / w2 - 1) < 1e-6
    # It lies near 2 c_p / h for the grid-scale in-plane mode.
    assert 0.8 < w1 * g.hx / (2 * STEEL.in_plane_speed) < 1.5
    s1, s2 = assemble(Plate.uniform(g, STEEL, 1e-3)), assemble(Plate.uniform(g, STEEL, 2e-3))
    c = s1.n // 2
    assert abs(s2.K[c, c] / s1.K[c, c] - 8) < 1e-9
    assert abs(s2.mass[c] / s1.mass[c] - 2) < 1e-12


def test_the_band_is_narrow_whichever_axis_is_longer():
    wide = assemble(Plate.uniform(make_grid(0.4, 0.1, 0.01), BRONZE, 1e-3))
    tall = assemble(Plate.uniform(make_grid(0.1, 0.4, 0.01), BRONZE, 1e-3))
    assert wide.bandwidth == tall.bandwidth <= 2 * 11 + 2